Trie ও Radix Tree — prefix দিয়ে খোঁজা

class TrieNode:
    __slots__ = ('children', 'is_word')

    def __init__(self):
        self.children = {}
        self.is_word = False


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        n = self.root
        for ch in word:
            n = n.children.setdefault(ch, TrieNode())
        n.is_word = True

    def _walk(self, prefix):
        n = self.root
        for ch in prefix:
            n = n.children.get(ch)
            if n is None:
                return None
        return n

    def search(self, word):
        n = self._walk(word)
        return n is not None and n.is_word

    def startswith(self, prefix):
        return self._walk(prefix) is not None

    def autocomplete(self, prefix):
        n = self._walk(prefix)
        if n is None:
            return []
        out = []
        def dfs(node, s):
            if node.is_word:
                out.append(s)
            for ch, kid in node.children.items():
                dfs(kid, s + ch)
        dfs(n, prefix)
        return out

    def count_nodes(self):
        def rec(n):
            return 1 + sum(rec(k) for k in n.children.values())
        return rec(self.root)


class RadixNode:
    __slots__ = ('edges', 'is_word')     # first-char -> (label, node)

    def __init__(self):
        self.edges = {}
        self.is_word = False


def build_radix(tnode):
    """trie-subtree → radix-subtree; single-child non-word chain merge"""
    rn = RadixNode()
    rn.is_word = tnode.is_word
    for ch, kid in tnode.children.items():
        label = ch
        while len(kid.children) == 1 and not kid.is_word:
            ((c2, kid),) = kid.children.items()     # একমাত্র child গোটাও
            label += c2
        rn.edges[label[0]] = (label, build_radix(kid))
    return rn


def count_radix(rn):
    return 1 + sum(count_radix(kid) for _, kid in rn.edges.values())


def radix_search(root, word):
    n, i = root, 0
    while i \< len(word):
        found = False
        for label, kid in n.edges.values():
            if word.startswith(label, i):
                i += len(label); n = kid; found = True; break
        if not found:
            return False
    return n.is_word


if __name__ == '__main__':
    words = ['car', 'card', 'care', 'cart', 'cat', 'dog']
    t = Trie()
    for w in words:
        t.insert(w)

    print('trie nodes :', t.count_nodes())
    assert sorted(t.autocomplete('ca')) == ['car', 'card', 'care', 'cart', 'cat']

    r = build_radix(t.root)
    print('radix nodes:', count_radix(r))
    for w in words:
        assert radix_search(r, w), f'{w} হারাল!'
    assert not radix_search(r, 'cars')
    assert not radix_search(r, 'do')
    print('radix query-ও একই উত্তর — semantics intact')
    print(f'savings: {t.count_nodes()} → {count_radix(r)} nodes')

প্রত্যাশিত output:

trie nodes : 11
radix nodes: 8
radix query-ও একই উত্তর — semantics intact
savings: 11 → 8 nodes

example-এর hand-count-এর সাথে মিলল (১১→৮)? নিয়ম আয়ত্তে। (ছোট set-এ modest savings — prefix-rich বড় dictionary-তে ৫০-৬০%, experiment-এ দেখেছেন।)

নিজে বাড়ান: (১) delete(word) + নতুন chain re-merge (বা lazy flag); (২) count_prefix(p) — node-এ word-count field রেখে O(|p|); (৩) Patricia-step: label-এর বদলে bit-skip-count কেন router-এ সুবিধা — ৫ লাইনে; (৪) বড় wordlist-এ depth-distribution মাপুন — max-depth নামে কি না।