Radiation search (Electronic station) - N*Nのマスをグループ分け

どんな問題？

Radiation search
http://www.checkio.org/mission/radiation-search/

N*Nのマスに、それぞれNo.1～5の数字が書いてある。同じ数字が縦横につながっている部分をグループとして、もっとも大きなグループのサイズ（マスの数）と数字を求めよ。
3 <= N <= 10 とする。

答えは、サイズと数字のリストで返す。No.3が10個のグループなら、 [10, 3]

なんか問題がわかりにくいけど、例えば、

1	1	2	2
1	1	2	1
1	1	2	1
3	3	1	1

この場合だと、

• No.1が6個のグループ
• No.2が4個のグループ
• No.3が2個のグループ
• No.1が4個のグループ

の4グループあって、最大のグループは、No.1が6個のグループだから、答えは [6, 1] となる。
グループは縦と横にはつながるが、斜めにはつながらない。単純に1～5のマスの数を数えればいいという問題ではないので注意。

例題：

checkio([
 [1,2,3,4,5],
 [1,1,1,2,3],
 [1,1,1,2,2],
 [1,2,2,2,1],
 [1,1,1,1,1]
])==[14,1]

どうやって解く？

まず、同じ番号のグループを探索する部分を考えてみる。開始マスから上下左右が同じ数字ならさらに探索、別の数字ならストップする。なんか、先日の迷路探索にも似ている。探索候補マスをスタックに入れていけばいいか。

１つのグループの探索が完了すれば、次は未探索のマスから開始して別のグループを探索する。こうやって全てのグループをリストアップし、最大のグループを求めればよい。

まず、マトリックスからグループを探索する部分を書いてみた。調査候補マスをスタックに入れていき、スタックが空になればループを抜ける。

import itertools

def get_group(matrix, start, visited):

    def get_number(p):
        # 座標pの数字を取得
        return matrix[p[1]][p[0]]

    def in_matrix(p):
        # 座標pがマトリックスの範囲内ならTrue
        n = len(matrix)
        return 0 <= p[0] < n and 0 <= p[1] < n

    checked = visited + [start]  # 調査済みフラグ
    stack = [start]  # 次に調査する座標（スタック）
    number = get_number(start)

    pts = []
    while len(stack) > 0:
        pt = stack.pop()
        pts.append(pt)
        # 上下左右のマスを調査
        for d in ((0, -1), (0, 1), (-1, 0), (1, 0)):  
            new_pt = pt[0] + d[0], pt[1] + d[1]
            if in_matrix(new_pt) and new_pt not in checked:
                checked.append(new_pt)
                if get_number(new_pt) == number:
                    stack.append(new_pt)

    return len(pts), number, pts

start座標が開始位置で、startと同じ数字のグループを探索して、(サイズ, 数字, 座標リスト) のタプルを返す。
引数のvisitedは、調査しなくて良い座標のリスト。次回探索時に、今回のグループの座標リストを指定しておけば、無駄な調査が減らせるので、少しは高速化できるはず。

メイン部分では、グループの探索と、未探索座標の取得を繰り返して、すべてのグループをリストアップする。

def checkio(matrix):
    def get_unvisited_pt():
        # 未探索の座標を取得
        all_pts = itertools.product(range(len(matrix)), repeat=2)
        unvisited = [p for p in all_pts if p not in visited]
        return unvisited[0] if unvisited else None

    pt = (0, 0)
    groups = []
    visited = []
    while pt is not None:
        size, number, pts = get_group(matrix, pt, visited)
        groups.append([size, number])
        visited += pts
        pt = get_unvisited_pt()

    return max(groups, key=lambda g: g[0])

一応完成したが、どうも無駄な探索が多いような気がする。
メイン部分checkio()のvisitedと、get_group()のcheckedは、どちらも探索済みフラグだから、まとめられないものだろうか。
あれこれ修正してみたが、結局あまり変わらなかったので、このまま公開する。

他の人の答え

あ、

return max(groups, key=lambda g: g[0])

この key= の部分は不要だった…。

さて、他の人の答えを見ると、スタックではなく再帰で書いている人が多かった。

def checkio(l):
    n = len(l)
    flags = [[0 for i in xrange(n)] for i in xrange(n)]
    def search(x, y, value):
 
        if not (0 <= x < n) or not (0 <= y < n):
            return 0
        if flags[x][y]:
            return 0
        current = l[x][y]
        if current == value:
            flags[x][y] = 1
            res = 1
            res += search(x-1, y, current)
            res += search(x+1, y, current)
            res += search(x, y-1, current)
            res += search(x, y+1, current)
            return res
        return 0
 
    res = max([search(x, y, l[x][y]), l[x][y]] for x in range(n) for y in range(n))
    return res

http://www.checkio.org/mission/radiation-search/publications/yumax/python-27/first/

これなんて、とてもシンプルでわかりやすい。
すべての座標についてグループ探索をしているので、すごく無駄のように見えるけど、そこはflagsで探索済みのマスは再調査しないから大丈夫。
なるほど。