trie.py (d86250a) - dost

trie.py

import os
import sys

def minima(options, key) -> set:
    xy = {o: key(o) for o in options}
    my = min(xy.values())
    return {x for (x, y) in xy.items() if y == my}

class RadixTree:
    def __init__(self, label: str = None, parent: "RadixTree" = None):
        self.label = label or ""
        self.children = set()
        self.parent = parent
        self.weight = 1 if label is not None else 0
        self.depth = parent.depth + 1 if parent else 0
        self.height = 1
        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):]
        match common:
            case self.label:
                if self.children:
                    for child in self.children:
                        if child.label:
                            ah = child.add(suffix)
                            if ah:
                                nh = max(self.height, 1 + ah)
                                break
                    else:
                        self.children.add(RadixTree(suffix, self))
                        nh = self.height
                else:
                    self.children = {RadixTree(suffix, self), RadixTree("", self)}
                    nh = 2
            case "":
                nh = None
            case c:
                rem = self.label[len(common):]
                nc = RadixTree(rem, self)
                nc.weight = self.weight
                nc.height = self.height
                nc.usage = self.usage
                nc.children = self.children
                for child in nc.children:
                    child.parent = nc
                self.label = c
                self.children = {RadixTree(suffix, self), nc}
                nh = 1 + self.height
        if nh:
            self.weight += 1
            self.height = nh
            return nh

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 method into a basic tool