trie.py (3d916e3) - dost

trie.py

from collections import defaultdict
import re
import sys
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: str, child: "Trie") -> tuple:
    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.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):
        if word:
            self.children[word[0]].add(word[1:], weight)
        else:
            self.weight += weight

def deep_weight(self) -> int:
        return self.weight + sum(c.deep_weight() for c in self.children.values())

def deep_usage(self) -> int:
        return self.usage + sum(c.deep_usage() for c in self.children.values())

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

def use(self):
        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
            leaf.use()

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

Remake trie.py using tokens instead of substrings