Skip to content

Instantly share code, notes, and snippets.

@cablehead
Last active October 6, 2020 11:10
Show Gist options
  • Save cablehead/8341c75854fe2d4c6d08 to your computer and use it in GitHub Desktop.
Save cablehead/8341c75854fe2d4c6d08 to your computer and use it in GitHub Desktop.
Iterate a 2d matrix using a spiral pattern
import unittest
import copy
from itertools import cycle
def row(matrix, n):
# returns row n of a 2d matrix
return matrix.pop(n)
def col(matrix, n):
# returns column n of a 2d matrix
return [row.pop(n) for row in matrix]
def rev(f):
# return a function such that rev(f) equals the reverse of the results of f
def mapped(*a, **kw):
return reversed(f(*a, **kw))
return mapped
def spiral(matrix):
# avoid mutating our caller's matrix
matrix = copy.deepcopy(matrix)
slice = cycle([row, col, rev(row), rev(col)])
side = cycle([0, -1, -1, 0])
# alternate between slicing rows and columns, and first and final side until
# the matrix has been consumed
while matrix:
for item in slice.next()(matrix, side.next()):
yield item
class TestSpiral(unittest.TestCase):
def test_input1(self):
input1 = [
[0, 1, 2, 3],
[4, 5, 6, 7],
[8, 9, 10, 11],
[12, 13, 14, 15], ]
output1 = [0, 1, 2, 3, 7, 11, 15, 14, 13, 12, 8, 4, 5, 6, 10, 9]
self.assertEqual(list(spiral(input1)), output1)
def test_input2(self):
input2 = [
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24], ]
output2 = [
0, 1, 2, 3, 4, 9, 14, 19, 24, 23, 22, 21, 20, 15, 10, 5, 6, 7, 8,
13, 18, 17, 16, 11, 12,]
self.assertEqual(list(spiral(input2)), output2)
if __name__ == '__main__':
unittest.main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment