Trie (prefix tree): Python-এ autocomplete, test ও সময়সহ

Trie (উচ্চারণ "ট্রাই", retrieval থেকে), অর্থাৎ prefix tree, শব্দকে অক্ষর ধরে ধরে রাখে, যাতে একই শুরুর শব্দগুলো একই পথ ভাগ করে নেয়। Autocomplete, বানান পরীক্ষক, অভিধান খোঁজা ও IP routing table-এর জন্য এটি ক্লাসিক structure, আর interview-তেও প্রায়ই আসে ("trie বানান", "word search II", "সবচেয়ে লম্বা সাধারণ prefix")। নিচের কোড, test ও সময় প্রকাশের আগে Python 3-এ চালানো হয়েছে।

ধারণাটি

car, card, care, careful ও cat রাখুন। list-এ c-a-r অক্ষর চারবার লেখা হয়। trie-তে একবারই:

c → a → r পথ চারটি শব্দের জন্য একই,
r-এর node-এ চিহ্ন থাকে "এখানে একটি শব্দ শেষ" (কারণ car নিজেও শব্দ),
তার নিচে d এবং e → f → u → l লম্বা শব্দগুলোতে গেছে,
a থেকে c → a → t আলাদা ডালে গেছে।

car দিয়ে শুরু সব শব্দ খুঁজতে তিন ধাপ নামুন, তারপর সেই node-এর নিচের সবকিছু তুলুন। dog বা dot কখনো দেখতেই হয় না।

কাজ	খরচ	মন্তব্য
L দৈর্ঘ্যের শব্দ insert	O(L)	প্রতি অক্ষরে এক ধাপ
সঠিক শব্দ খোঁজা	O(L)	নামুন, তারপর শেষ-চিহ্ন দেখুন
"এই দিয়ে শুরু কোনো শব্দ আছে?"	O(L)	শুধু নামুন
Autocomplete	O(L + উত্তরের আকার)	prefix-এ নামুন, তারপর তার নিচের অংশ ঘুরুন

খরচ নির্ভর করে শব্দের দৈর্ঘ্যের ওপর, কতগুলো শব্দ রাখা আছে তার ওপর নয়। এটাই এর আসল আকর্ষণ।

একটি চালানো যায় এমন উদাহরণ

trie.py নামে সংরক্ষণ করুন:

import random
import string
import time


class Trie:
    """A prefix tree: each node maps a character to a child node."""

    def __init__(self):
        self.root = {}
        self.size = 0

    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.setdefault(ch, {})
        if "$" not in node:  # "$" marks the end of a stored word
            node["$"] = True
            self.size += 1

    def _walk(self, text):
        node = self.root
        for ch in text:
            node = node.get(ch)
            if node is None:
                return None
        return node

    def contains(self, word):
        node = self._walk(word)
        return node is not None and "$" in node

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

    def autocomplete(self, prefix, limit=5):
        """Return up to `limit` stored words that start with prefix, in alphabetical order."""
        node = self._walk(prefix)
        out = []
        if node is None:
            return out

        def visit(n, path):
            if len(out) >= limit:
                return
            if "$" in n:
                out.append(prefix + "".join(path))
            for ch in sorted(k for k in n if k != "$"):
                visit(n[ch], path + [ch])

        visit(node, [])
        return out


def list_prefix_scan(words, prefix, limit=5):
    """The simple way: look at every word."""
    return sorted(w for w in words if w.startswith(prefix))[:limit]


if __name__ == "__main__":
    t = Trie()
    for w in ["car", "card", "care", "careful", "cat", "dog", "do", "dot"]:
        t.insert(w)
    print("contains('car')      :", t.contains("car"))
    print("contains('ca')       :", t.contains("ca"), "(a prefix, not a stored word)")
    print("starts_with('ca')    :", t.starts_with("ca"))
    print("autocomplete('car')  :", t.autocomplete("car"))
    print("autocomplete('do', 2):", t.autocomplete("do", 2))
    print("autocomplete('x')    :", t.autocomplete("x"))

    rnd = random.Random(2)
    words = {"".join(rnd.choice(string.ascii_lowercase) for _ in range(rnd.randint(4, 10))) for _ in range(150_000)}
    big = Trie()
    t0 = time.perf_counter()
    for w in words:
        big.insert(w)
    build = (time.perf_counter() - t0) * 1000
    wl = list(words)
    prefix = "qu"
    n = 200
    t0 = time.perf_counter()
    for _ in range(n):
        a = big.autocomplete(prefix)
    t1 = (time.perf_counter() - t0) * 1000 / n
    t0 = time.perf_counter()
    for _ in range(n):
        b = list_prefix_scan(wl, prefix)
    t2 = (time.perf_counter() - t0) * 1000 / n
    print(f"\n{len(words):,} words, build trie: {build:.0f} ms")
    print(f"same answer for prefix '{prefix}': {a == b}")
    print(f"  trie autocomplete   {t1:8.3f} ms per query")
    print(f"  scan every word     {t2:8.3f} ms per query")

python3 trie.py-এর ফল (সঠিক সময় মেশিনভেদে বদলায়):

contains('car')      : True
contains('ca')       : False (a prefix, not a stored word)
starts_with('ca')    : True
autocomplete('car')  : ['car', 'card', 'care', 'careful']
autocomplete('do', 2): ['do', 'dog']
autocomplete('x')    : []

149,469 words, build trie: 716 ms
same answer for prefix 'qu': True
  trie autocomplete      0.026 ms per query
  scan every word       13.679 ms per query

কোড কীভাবে কাজ করে

প্রতিটি node একটি dictionary, যা একটি অক্ষর থেকে পরের node-এ যায়। বিশেষ key "$" একটি রাখা শব্দের শেষ বোঝায়। এভাবেই contains('ca') False কিন্তু starts_with('ca') True: পথ আছে, কিন্তু সেখানে কোনো শব্দ শেষ হয় না।
insert setdefault দিয়ে node বানিয়ে নামে, তারপর শেষ-চিহ্ন বসায়। একই শব্দ আবার দিলে size বদলায় না।
autocomplete আগে prefix-এর node-এ যায়। তারপর depth-first ভ্রমণে অক্ষর সাজানো ক্রমে সন্তানদের ঘোরে, তাই ফল বর্ণানুক্রমে আসে আর সীমায় থামে।
list_prefix_scan সহজ বিকল্প: প্রতিটি শব্দে startswith পরীক্ষা। এটি আছে যাতে trie-কে একটি নির্ভরযোগ্য উত্তরের সঙ্গে মেলানো যায়।

সময় পড়া

প্রায় ১,৫০,০০০ random শব্দ আর prefix qu নিয়ে:

trie বানাতে লাগল ৭১৬ ms (একবারের খরচ)।
trie autocomplete লাগল প্রতি query-তে ০.০২৬ ms।
প্রতিটি শব্দ scan করতে প্রতি query-তে ১৩.৭ ms, প্রায় ৫০০ গুণ ধীর।

scan শব্দসংখ্যার সঙ্গে ধীর হয়, trie query হয় না। শুধু একটি query-র জন্য scan জেতে, কারণ বানানোর ধাপ লাগে না। একই শব্দভান্ডারে অনেক query, যেমন search box, সেখানে trie অনায়াসে জেতে।

Test

test_trie.py নামে সংরক্ষণ করুন:

import random
import string
import unittest

from trie import Trie, list_prefix_scan


class TrieTests(unittest.TestCase):
    def test_contains_vs_prefix(self):
        t = Trie()
        t.insert("car")
        self.assertTrue(t.contains("car"))
        self.assertFalse(t.contains("ca"))
        self.assertTrue(t.starts_with("ca"))
        self.assertFalse(t.starts_with("cb"))

    def test_duplicates_counted_once(self):
        t = Trie()
        for w in ["a", "a", "ab"]:
            t.insert(w)
        self.assertEqual(t.size, 2)

    def test_empty_prefix_lists_everything_sorted(self):
        t = Trie()
        for w in ["b", "a", "c"]:
            t.insert(w)
        self.assertEqual(t.autocomplete("", 10), ["a", "b", "c"])

    def test_matches_simple_scan_on_random_words(self):
        rnd = random.Random(4)
        words = {"".join(rnd.choice("abc") for _ in range(rnd.randint(1, 6))) for _ in range(300)}
        t = Trie()
        for w in words:
            t.insert(w)
        for prefix in ["", "a", "ab", "abc", "cc", "zzz"]:
            self.assertEqual(t.autocomplete(prefix, 7), list_prefix_scan(words, prefix, 7))

    def test_limit_respected(self):
        t = Trie()
        for w in ["aa", "ab", "ac", "ad"]:
            t.insert(w)
        self.assertEqual(len(t.autocomplete("a", 2)), 2)


if __name__ == "__main__":
    unittest.main()

$ python3 -m unittest
Ran 5 tests in 0.004s

OK

সবচেয়ে কাজের test হলো trie-কে সহজ scan-এর সঙ্গে random শব্দ ও নানা prefix-এ মেলানো, খালি prefix ও কোনো মিল নেই এমন prefix-সহ।

Trie নাকি অন্য কিছু?

স্থির বা ধীরে বদলানো শব্দসমষ্টিতে অনেক prefix query: trie।
শুধু সঠিক শব্দ খোঁজা: hash set সহজ ও দ্রুত। trie লাগে না।
সাজানো list ও binary search: prefix-এর পরিসরও O(L log n)-এ পাওয়া যায়, memory অনেক কম। data কম বদলালে এটি ভালো বিকল্প।
Memory: প্রতি node-এ dictionary দেওয়া সাধারণ trie বড় শব্দসমষ্টিতে অনেক memory নেয়। সংকুচিত রূপ (radix tree, DAWG) ও নির্দিষ্ট বর্ণমালার জন্য array এটি কমায়।

সাধারণ ভুল

শেষ-চিহ্ন ভুলে যাওয়া, ফলে card আছে বলেই car-ও আছে মনে হয়।
"prefix আছে" আর "শব্দ আছে" গুলিয়ে ফেলা।
মাত্র একটি query-র জন্য trie বানানো।
autocomplete-এর ফল এলোমেলো ক্রমে ফেরত দেওয়া।
লক্ষ লক্ষ শব্দে memory, বা বড়-ছোট হাতের অক্ষর ও accent উপেক্ষা করা (আগে text normalise করুন)।

মূল কথা: trie সাধারণ শুরুর অংশ ভাগ করে নেয়, তাই prefix খোঁজার খরচ হয় prefix-এর দৈর্ঘ্য আর উত্তরের আকার, অভিধানের আকার নয়। autocomplete ও prefix query-তে এটি ব্যবহার করুন, আর সহজ scan-এর সঙ্গে মিলিয়ে test করুন।

Interview প্রস্তুতি ও DSA সিরিজ, পর্ব ৩ (আরও algorithm): 3 নম্বর (মোট 5)। আগের লেখা: Heap ও priority queue: top-k, merge ও task scheduling.