Searching a 2D matrix for a consecutive sequence

Question

I've been solving problems on checkio (and trying to digest other's code) in order to improve my Python.

My main goals are to learn to write more idiomatic Python and to explore the language more fully ("batteries included" is no joke!). I'd like critiques on the following (working) code in terms of pythonic-ness and in terms of my choice of tools/solution method:

def checkio(matrix):
 def four_consecutive(group):
 for number in set(group):
 if str(number)*4 in ''.join(map(str,group)): return True
 return False
 def search(generators):
 for generator in generators:
 while True:
 try:
 if four_consecutive(next(generator)): return True
 except StopIteration: break
 return False
 horizontal = (row for row in matrix) 
 vertical = ([matrix[i][j] for i in range(len(matrix[j]))] for j in range(len(matrix)))
 diagonal1 = ([matrix[j+i][j] for j in range(len(matrix)-i)] for i in range(len(matrix)-3))
 diagonal2 = ([matrix[j][j+i] for j in range(len(matrix)-i)] for i in range(1,len(matrix)-3))
 diagonal3 = ([matrix[-1*(j+1)][j+i] for j in range(len(matrix)-i)] for i in range(len(matrix)-3))
 diagonal4 = ([matrix[-1*(j+i+1)][j] for j in range(len(matrix)-i)] for i in range(1,len(matrix)-3))
 if search((horizontal,vertical,diagonal1,diagonal2,diagonal3,diagonal4)): return True
 return False

Matrices in this problem are always square, contain integers 1-9, and the specification asks for a return value of true if 4 consecutive equal integers are in a line horizontally, vertically, or diagonally, and false otherwise. Core Python modules were available for import.

My main goal lately has been to get more comfortable with generators and list comprehensions, and to write shorter code.

Answer 1

Combining @rolfls windowing with generators is fairly simple:

import itertools
def consecutive(group):
 first, second = itertools.tee(group)
 second.next()
 for first, second in itertools.izip(first, second):
 if second != first + 1: return False
 return True
def iterate_submatrix(matrix, t, l):
 '''yield the horizontals and diagonals of 4x4 subsection of matrix starting at t(op), l(eft) as 4-tuples'''
 submat = [row[l:l+4] for row in matrix[t:t+4]]
 for r in submat: yield tuple(r) 
 for c in range (0,4): 
 yield tuple(r[c] for r in submat)
 yield tuple(submat[rc][rc] for rc in range (0,4))
 yield tuple(submat[rc][3-rc] for rc in range(0,4))
# sample usage:
for item in iterate_submatrix(test_matrix, 0,0):
 print item, consecutive(item)

There is probably some perf overhead in the generators here, but they do have the pleasant property of hiding several different selection styles behind a neat facade and also minimizing the index mongering. You could easily parameterize the windowing code to support larger or shorter sequences too.

• Hopping Mouse: Yes to minimizing index mongering, itertools, and avoiding string conversion. Learning itertools and functools are what I'm working on right now. Thanks!

Answer 2

The efficiency of this check must surely be suffering because of the final check.... create a String for each four-some, and compare it against a reference String. In general this will be slow, but, in particular, there can only be 9 reference strings, so why do you have to calculate it each time? An array with these 9 values would be simple to create, and then you can just say:

if reference[number] in ''.join(map(str,group)): return True

I am not suggesting that this is a good solution, just a better/faster solution. The best/fastest solution would not create two string values for comparison, but would leave the values in their native integer format.

Secondly, why do you have 4 diagonals? You should be able to get by with just 2 (since there are only two of them). I have not been able to get my head around this problem quickly enough to understand whether there is a particular use-case requiring the 4, but at face value, this means that you are double-checking your diagonals.

While I understand that you are trying to get a handle on generators and comprehensions, a more simplistic approach to this problem would be more efficient, and probably more readable.

The standard HPC (high-performance-computing) method for this type of problem is to create a logical box that spans the problems space. In this case, the box would be 4 by 4. You then create a cursor that traverses the problem space, where the bottom-right member of the box is on the cursor.

You then test that box for a vertical line above, the cursor (if there is space), a horizontal-line to the left, and the two diagonals in the box.

This type of approach is easy to multi-thread or distribute. There are ways to make the memory access more efficient as well by splitting the problem on clever boundaries when parallelizing.

But, what you want to avoid is generating values for each movement of the box....

• Hopping Mouse: Good point about the reference array. I think that my method of getting sublists with the six generators is similar to the HPC method you detailed. The 'cursor' is the outer (mostly 'i') variable, and it walks along an 'edge' of the matrix generating 'lines' as lists. There are two generators per diagonal 'direction' walking the cursor down the 'side' and across the 'top' of the matrix to get all the diagonals. So diagonals 1 and 2 handle 'down-right' while 3 and 4 handle 'down-left'. Each one contains unique data.