trie.py (d6c605c) - dost

trie.py

import os
import sys

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

class RadixTree:
    def __init__(self, label: str = None, parent: "RadixTree" = None):
        self.label = label or ""
        self.children: set[RadixTree] = set()
        self.parent = parent
        self.weight = 1 if label is not None else 0
        self.depth = parent.depth + 1 if parent else 0
        self.usage = 0

def __str__(self):
        s = f"{self.label or "∅"}"
        if self.children:
            s = f"({s} {" ".join(str(x) for x in self.children)})"
        return s

def add(self, word: str):
        common = os.path.commonprefix([self.label, word])
        suffix = word[len(common):]
        if common == self.label:
            for child in self.children:
                if child.label:
                    if child.add(suffix):
                        break
            else:
                if not self.children:
                    self.children.add(RadixTree("", self))
                self.children.add(RadixTree(suffix, self))
        elif common:
            remainder = self.label[len(common):]
            split = RadixTree(remainder, self)
            split.weight = self.weight
            split.usage = self.usage
            split.children = self.children
            for child in split.children:
                child.parent = split
            self.label = common
            self.children = {RadixTree(suffix, self), split}
        if common or not self.label:
            self.weight += 1
            return True

def full(self) -> str:
        s = self.label
        if self.parent:
            s = self.parent.full() + s
        return s

def del_empty(self):
        for child in list(self.children):
            if not child.label:
                if child.children:
                    child.del_empty()
                else:
                    self.children.remove(child)

def leaves(self) -> set["RadixTree"]:
        if self.children:
            return {l for c in self.children for l in c.leaves()}
        else:
            return {self}

def use(self):
        assert self.usage < self.weight
        self.usage += 1
        if self.parent:
            self.parent.use()

def best(self) -> "RadixTree":
        if self.children:
            branches = minima(self.children, lambda c: (c.usage / c.weight, -(c.weight // c.depth)))
            return min((b.best() for b in branches), key=lambda b: (len(b.full()), b.full()))
        else:
            return self

t = RadixTree()
words = {w.strip() for w in sys.stdin}
for w in words:
    t.add(w.strip())
if "" not in words:
    t.del_empty()
leaves = {l.full() for l in t.leaves()}
assert leaves <= words
assert words <= leaves
while t.usage < t.weight:
    b = t.best()
    b.use()
    print(b.full())

dost

Refine trie.py by advice of GPT