Trie Applications字典樹應用

from collections import deque


class AhoCorasick:
    """多模式比對：建表 O(關鍵字總長)，掃描 O(n + 回報次數)"""

    def __init__(self, words):
        self.goto = [{}]                        # 每個節點的子節點：字元 → 節點編號
        self.fail = [0]
        self.out = [[]]                         # 走到這個節點時要回報的關鍵字
        for w in words:
            node = 0
            for ch in w:
                if ch not in self.goto[node]:
                    self.goto[node][ch] = len(self.goto)
                    self.goto.append({})
                    self.fail.append(0)
                    self.out.append([])
                node = self.goto[node][ch]
            self.out[node].append(w)
        q = deque(self.goto[0].values())        # 第一層的失敗連結都指向根
        while q:                                # BFS：失敗連結一定指向更淺、已經算好的節點
            u = q.popleft()
            for ch, v in self.goto[u].items():
                q.append(v)
                f = self.fail[u]
                while f and ch not in self.goto[f]:
                    f = self.fail[f]
                self.fail[v] = self.goto[f].get(ch, 0)
                self.out[v] = self.out[v] + self.out[self.fail[v]]   # 後綴也是關鍵字時一起回報

    def search(self, text):
        res, node = [], 0
        for i, ch in enumerate(text):
            while node and ch not in self.goto[node]:
                node = self.fail[node]          # 走不下去就沿失敗連結往回跳
            node = self.goto[node].get(ch, 0)
            for w in self.out[node]:
                res.append((i - len(w) + 1, w))
        return res


def find_words(board, words):
    """Word Search II：DFS 時帶著字典樹節點，前綴不存在就立刻剪枝"""
    root = {}
    for w in words:
        node = root
        for ch in w:
            node = node.setdefault(ch, {})
        node["$"] = w                           # 結尾標記直接存整個單字
    rows, cols = len(board), len(board[0])
    found = []

    def dfs(r, c, parent):
        ch = board[r][c]
        node = parent.get(ch)
        if node is None:                        # 沒有任何單字以這個前綴開頭
            return
        if "$" in node:
            found.append(node.pop("$"))         # 收下並清掉，避免重複回報
        board[r][c] = "#"
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and board[nr][nc] != "#":
                dfs(nr, nc, node)
        board[r][c] = ch
        if not node:
            parent.pop(ch)                      # 這段前綴底下的單字都找到了，整段剪掉

    for r in range(rows):
        for c in range(cols):
            dfs(r, c, root)
    return sorted(found)


def max_xor_pair(nums, bits=31):
    """把整數的二進位從最高位開始當成字串插進字典樹；查詢時每一位都盡量走相反的位元"""
    root, best = {}, 0
    for x in nums:
        node = root
        for b in range(bits - 1, -1, -1):
            node = node.setdefault((x >> b) & 1, {})
        node, cur = root, 0
        for b in range(bits - 1, -1, -1):
            want = 1 - ((x >> b) & 1)
            if want in node:
                cur |= 1 << b
                node = node[want]
            else:
                node = node[1 - want]
        best = max(best, cur)
    return best


if __name__ == "__main__":
    ac = AhoCorasick(["he", "she", "his", "hers"])
    print(ac.search("ushers"))                  # [(1, 'she'), (2, 'he'), (2, 'hers')]
    board = [list("oaan"), list("etae"), list("ihkr"), list("iflv")]
    print(find_words(board, ["oath", "pea", "eat", "rain"]))   # ['eat', 'oath']
    print(max_xor_pair([3, 10, 5, 25, 2, 8]))   # 28：5 XOR 25