trie.py (fd1af19) - dost

trie.py

from collections import defaultdict
import re
import sys
import time
from typing import Generator

def minima(options, key):
    min_val = None
    result = set()
    for o in options:
        k = key(o)
        if min_val is None or k < min_val:
            min_val = k
            result = {o}
        elif k == min_val:
            result.add(o)
    return result

def cost(token_child: tuple[str, "Trie"]) -> tuple:
    token, child = token_child
    weight = child.deep_weight
    return child.deep_usage / weight, -weight, len(token), token

class Trie:
    def __init__(self, weight: int = 0):
        self.children: defaultdict[str, Trie] = defaultdict(Trie)
        self.weight: int = weight
        self.deep_weight: int = weight
        self.usage: int = 0
        self.deep_usage: int = 0

def __str__(self) -> str:
        s = f"[{self.usage}/{self.weight}]"
        if self.children:
            s += f"({" ".join('"' + k + '"' + str(v) for (k, v) in self.children.items())})"
        return s

def add(self, word: list[str], weight: int = 1):
        self.deep_weight += weight
        if word:
            self.children[word[0]].add(word[1:], weight)
        else:
            self.weight += weight

def best(self) -> tuple[list[str], "Trie"]:
        if self.weight > self.usage:
            return [], self
        elif self.children:
            token, child = min(self.children.items(), key=cost)
            suffix, leaf = child.best()
            if not leaf:
                return None, None
            return [token] + suffix, leaf
        else:
            return None, None

def use(self, word: list[str]):
        self.deep_usage += 1
        if word:
            self.children[word[0]].use(word[1:])
        else:
            assert self.usage < self.weight
            self.usage += 1

def traverse(self) -> Generator[tuple[list[str], "Trie"], None, None]:
        while True:
            tokens, leaf = self.best()
            if not leaf:
                assert self.deep_usage == self.deep_weight
                break
            yield tokens
            self.use(tokens)

re_flags = re.MULTILINE | re.DOTALL
token_pattern = "[-_ \t]*[^-_\\s]+"
entry_pattern = re.compile(f"({token_pattern})+", re_flags)
token_pattern = re.compile(token_pattern, re_flags)
words = [re.findall(token_pattern, m.group(0)) for m in re.finditer(entry_pattern, sys.stdin.read())]
t = Trie()
for w in words:
    t.add(w)
for w in t.traverse():
    print("".join(w))

dost

Make trie method way faster via "caching"