745. 前缀和后缀搜索

exiaohu 于 2022-07-14 发布

题目链接:745. 前缀和后缀搜索

字典树的变体,针对这个问题,有两个重要改进:

  1. 字典树中每个节点的 key 是一个前、后缀的元组;
  2. 在每个节点都保存一个以本节点为根的子树的叶子节点对应的单词的索引的最大值。
from itertools import zip_longest
from typing import List


class WordFilter:

    def __init__(self, words: List[str]):
        trie = {}
        for k, word in enumerate(words):
            cur, drow = trie, word[::-1]
            for i, (p, s) in enumerate(zip(word, drow)):
                tmp = cur
                for l, r in zip_longest(word[i:], '', fillvalue='#'):
                    if (l, r) not in tmp:
                        tmp[(l, r)] = {}
                    tmp = tmp.get((l, r))
                    tmp[('#', '#')] = k

                tmp = cur
                for l, r in zip_longest('', drow[i:], fillvalue='#'):
                    if (l, r) not in tmp:
                        tmp[(l, r)] = {}
                    tmp = tmp.get((l, r))
                    tmp[('#', '#')] = k

                if (p, s) not in cur:
                    cur[p, s] = {}
                cur = cur.get((p, s))
                cur[('#', '#')] = k

        self.trie, self.words = trie, {word: i for i, word in enumerate(words)}

    def f(self, pref: str, suff: str) -> int:
        trie = self.trie
        for p, s in zip_longest(pref, suff[::-1], fillvalue='#'):
            if not trie:
                return -1

            trie = trie.get((p, s), {})

        return trie.get(('#', '#'), -1)