【Python】単語検索II（Word Search II）をTrieとDFSで効率的に解く方法

はじめに

本記事では、2次元の文字ボードと単語リストが与えられたとき、ボード上から辞書内のすべての単語を見つけ出す「単語検索II（Word Search II）」という定番アルゴリズム問題を、Pythonで解いていきます。

この問題のルールは次の通りです。

• 各単語は、上下左右に隣接するセルの文字を順番につなげて構成する必要がある
• 同じセル（マス）を1つの単語の中で複数回使用することはできない

例えば、後述のような4×4のボードと単語リスト ["oath", "pea", "tea", "rain"] が入力として与えられた場合を考えてみましょう。

解法のアプローチ：Trie（トライ木）＋ DFSバックトラッキング

この問題を効率よく解くには、Trie（トライ木）と深さ優先探索（DFS）を組み合わせるのが定石です。単語ごとに毎回ボード全体を走査すると計算量が膨大になりますが、あらかじめ単語リストをTrieに登録しておけば、探索中に「その文字から始まる未発見の候補単語がまだ残っているか」を即座に判定でき、無駄な分岐を大幅に削減できます。

アルゴリズムの手順

• 結果を格納する配列 result を作成する
• solve() メソッドを定義する（引数は board、d、i、j、s）
• i または j がボードの行・列の範囲外であれば return する
• l := board[i][j]
• l が d に存在する場合：
  • d := d[l] とし、s に l を連結する
  • d に "#" が存在し、かつ d["#"] が真であれば：
    • s を result に追加する
    • d["#"] := 0 に設定し、同じ単語の重複検出を防ぐ
  • board[i][j] := "*" とし、同一セルの再利用を一時的に禁止する
  • 下方向：i+1 が範囲内かつ board[i+1][j] が d にあれば solve(board, d, i+1, j, s) を呼び出す
  • 右方向：j+1 が範囲内かつ board[i][j+1] が d にあれば solve(board, d, i, j+1, s) を呼び出す
  • 上方向：i-1 が範囲内かつ board[i-1][j] が d にあれば solve(board, d, i-1, j, s) を呼び出す
  • 左方向：j-1 が範囲内かつ board[i][j-1] が d にあれば solve(board, d, i, j-1, s) を呼び出す
  • board[i][j] := l に復元する（バックトラック）
• insert() メソッドを定義する（引数は word と辞書 t）
• current := t
• word の各文字 i について：
  • i が current に存在しなければ current[i] := 新しいマップ
  • current := current[i]
• current["#"] := 1（単語の終端マーク）

メイン処理の流れ

• 空のマップ t を作成する
• words の各 word に対して insert(word, t) を呼び出す
• board のすべてのセル (i, j) に対して solve(board, t, i, j) を呼び出す
• result を返す

実装例

それでは、実際のコードを見ながら理解を深めていきましょう。

class Solution(object):
    def findWords(self, board, words):
        self.result = []
        t = {}
        for word in words:
            self.insert(word,t)
        for i in range(len(board)):
            for j in range(len(board[0])):
                self.solve(board,t,i,j)
        return self.result
    def solve(self,board,d,i,j,s=""):
        if i<0 or j<0 or i>=len(board) or j>=(len(board[0])):
            return
        l = board[i][j]
        if l in d:
            d = d[l]
            s+=l
            if "#" in d and d['#']:
                self.result.append(s)
                d['#'] = 0
            board[i][j] = '*'
            if i+1<len(board) and board[i+1][j] in d :
                self.solve(board,d,i+1,j,s)
            if j+1 < len(board[0]) and board[i][j+1] in d:
                self.solve(board,d,i,j+1,s)
            if i-1>=0 and board[i-1][j] in d :
                self.solve(board,d,i-1,j,s)
            if j-1>=0 and board[i][j-1] in d :
                self.solve(board,d,i,j-1,s)
            board[i][j] = l
    def insert(self, word,t):
        current = t
        for i in word:
            if i not in current:
                current[i] = {}
            current =current[i]
        current['#']=1

ob = Solution()
print(ob.findWords([["o","a","a","n"],["e","t","e","a"],["i","h","k", "r"],["i","f","l","v"]],["oath","pea","tea","rain"]))

入力

[["o","a","a","n"],
["e","t","e","a"],
["i","h","k","r"],
["i","f","l","v"]],
["oath","pea","tea","rain"]

出力

['oath', 'tea']

実装のポイント解説

• Trieによる前処理： 単語リストをネストした辞書構造に変換することで、DFSの各ステップで「この枝はもう探索する価値がない」と即座に判断できます。
• 終端マーク "#" の活用： 単語の末尾を "#" で表現し、単語が見つかったらフラグを 0 に書き換えることで、同じ単語が結果に重複して追加されるのを防いでいます。
• バックトラッキング： 探索中は現在のセルを "*" で上書きして再利用を防ぎ、探索が終わったら元の文字に復元します。これにより、別の経路からの探索にも正しく対応できます。

このように、TrieとDFSを組み合わせた手法は、大規模なボードや多数の単語を扱う場合でも高いパフォーマンスを発揮する、非常に実用的なアプローチです。