Find Sequence (Electronic Station) - 連続した同じ数字を探せ

どんな問題？

Find Sequence
http://www.checkio.org/mission/find-sequence/

NxN (4<=N<=10) のマトリックスの中から、縦・横・斜めに4個以上の連続した同じ数字があるかどうかを調べる。
存在すればTrue、なければFalseを返す。

どうやって解く？

前のXs and Os (○×ゲーム)に似ているので、同じようなアプローチで考えてみる。

1. マトリックスから、縦、横、斜めの数字列を列挙していく。
2. 列挙された数字列で、同じ数字が4つ以上連続しているかどうかチェックする。

という2ステップでいけそうだ。
縦と横の数字列はいいとして、問題は斜めかな。例えば、5*5のマトリックスの場合だと、縦5本、横5本、斜め6本の数字列をチェックしないといけない。

def get_number(matrix, pt):
    return matrix[pt[1]][pt[0]]

def enum_lines(n):
    """
    N*Nのマトリックスの縦、横、斜めの並びの座標を返す
    list of list of (x,y)
    """
    def enum_line(pt, count, get_next_pt):
        #funcで点を移動しながら列挙する
        for i in range(count):
            if 0 <= pt[0] < n and 0 <= pt[1] < n:  # 範囲チェック
                yield pt
            pt = get_next_pt(pt)

    def get_line(pt, count, get_next_pt):
        # generatorをlist化
        return list(enum_line(pt, count, get_next_pt))

    # 縦
    for x in range(n):
        yield get_line((x, 0), n, lambda pt: (pt[0], pt[1] + 1))
    # 横
    for y in range(n):
        yield get_line((0, y), n, lambda pt: (pt[0] + 1, pt[1]))
    # 斜め
    for x in range(-(n - 2), n - 1):  # n=4, -2 to 3
        # 左上から右下へ
        yield get_line((x, 0), n, lambda pt: (pt[0] + 1, pt[1] + 1))
        # 左下から右上へ
        yield get_line((x, n - 1), n, lambda pt: (pt[0] + 1, pt[1] - 1))

def has_sequence(numbers, count):
    """
    numbersにcount個以上の連続した数字があるかどうかをチェック
    (例) has_sequence([1,2,3,3,3,3], 4) == True
    """
    for i in range(len(numbers) - count + 1):
        if numbers[i:i + count] == [numbers[i]] * count:
            return True

def checkio1(matrix):
    SEQ_LENGTH = 4

    # 縦横斜めのラインの座標をリスト化  (例) [[(0,0), (0,1), (0,2), (0,3)], [(1,0), (1,1), (1,2), (1,3)],...
    lines = [line for line in enum_lines(len(matrix)) if len(line) >= SEQ_LENGTH]  # 4個未満のラインは除外

    # 座標を数値に変換  (例) [[1, 2, 1, 1], [1, 1, 4, 1], [1, 3, 1, 6], [1, 7, 2, 5]]
    num_lines = [[get_number(matrix, pt) for pt in line] for line in lines if len(line) >= SEQ_LENGTH]

    # 各ラインで、4個以上の連続した数字があるかチェック
    return any(has_sequence(line, SEQ_LENGTH) for line in num_lines)

1. マトリックスから、縦横斜めのラインをリストアップする。 enum_lines()
2. ラインを数字列に変換する。
3. 数字列に連続した数字があるかチェックする。 has_sequence()

という3ステップになった。
ラインというのは、座標(x,y)のリストで、例えば、4*4のマトリックスなら、縦4本、横4本、斜め2本の10本のラインがある。斜めのラインは [(0,0),(1,1),(2,2),(3,3)] と [(0,3),(1,2),(2,1),(3,0)] となる。
座標の数字は get_number()で取得し、数字列に変換してから、連続した数字をチェックする、という流れ。

enum_lines()の斜めのラインをリストアップするあたりが、ちょっとわかりにくいかもしれない。
関数内のenum_line()は、開始点、点の数、隣の点への移動方法をラムダ式で指定すれば、ラインを取得できる。(0,0)から(5,0)への線なら、enum_line( (0,0),5,lambda pt: (pt[0]+1,pt[1]) ) となる。

やっぱり、わかりにくいな。もうちょっとシンプルに書き直してみた。

def enum_lines(n, min_length):
    def in_matrix(pt):
        return 0 <= pt[0] < n and 0 <= pt[1] < n

    def get_line(start, get_pt):
        return [pt for pt in [get_pt(start, i) for i in range(n)] if in_matrix(pt)]

    for k in range(n):
        # 横の線
        yield get_line((0, k), lambda pt, i: (pt[0] + i, pt[1]))
        # 縦の線
        yield get_line((k, 0), lambda pt, i: (pt[0], pt[1] + i))

    for x in range(-n + min_length, n - min_length + 1):
        # 左上から右下へ
        yield get_line((x, 0), lambda pt, i: (pt[0] + i, pt[1] + i))
        # 左下から右上へ
        yield get_line((x, n - 1), lambda pt, i: (pt[0] + i, pt[1] - i))

get_line()関数は、開始点と、開始点からの差分をラムダ式で指定するように変えた。(0,0)から(5,0)への線なら、get_line( (0,0),lambda pt,i: (pt[0]+i,pt[1]) ) となる。
min_lengthを指定して、その長さ未満のラインは取得しないようにした。これで、後からラインの長さチェックをする必要がなくなった。

全体をまとめるとこうなった。

def get_number(matrix, pt):
    return matrix[pt[1]][pt[0]]

def enum_lines(n, min_length):
    def in_matrix(pt):
        return 0 <= pt[0] < n and 0 <= pt[1] < n

    def get_line(start, get_pt):
        return [pt for pt in [get_pt(start, i) for i in range(n)] if in_matrix(pt)]

    for k in range(n):
        # 横の線
        yield get_line((0, k), lambda pt, i: (pt[0] + i, pt[1]))
        # 縦の線
        yield get_line((k, 0), lambda pt, i: (pt[0], pt[1] + i))

    for x in range(-n + min_length, n - min_length + 1):
        # 左上から右下へ
        yield get_line((x, 0), lambda pt, i: (pt[0] + i, pt[1] + i))
        # 左下から右上へ
        yield get_line((x, n - 1), lambda pt, i: (pt[0] + i, pt[1] - i))

def has_sequence(numbers, count):
    return any(numbers[i:i + count] == [numbers[i]] * count for i in range(len(numbers) - count + 1))

def checkio(matrix):
    SEQ_LENGTH = 4

    num_lines = ([get_number(matrix, pt) for pt in line] for line in enum_lines(len(matrix), SEQ_LENGTH))
    return any(has_sequence(numbers, SEQ_LENGTH) for numbers in num_lines)

もうちょっと改良できそうだけど、まあ、これでいいや。さくっと公開！
あれ？Clear カテゴリで公開してるの私だけ？？？