dost

@@ -1,5 +1,7 @@

-import os

+from collections import defaultdict

+import re

 import sys

+from typing import Generator

 def minima(options, key):

     min_val = None

@@ -13,90 +15,69 @@ def minima(options, key):

             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 cost(token: str, child: "Trie") -> tuple:

+    weight = child.deep_weight()

+    return child.deep_usage() / weight, -weight, len(token), token

-    def __str__(self):

-        s = f"{self.label or "∅"}"

+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"({s} {" ".join(str(x) for x in self.children)})"

+            s += f"({" ".join('"' + k + '"' + str(v) for (k, v) in self.children.items())})"

         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

+    def add(self, word: list[str], weight: int = 1):

+        if word:

+            self.children[word[0]].add(word[1:], weight)

         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

+            self.weight += weight

-    def full(self) -> str:

-        s = self.label

-        if self.parent:

-            s = self.parent.full() + s

-        return s

+    def deep_weight(self) -> int:

+        return self.weight + sum(c.deep_weight() for c in self.children.values())

-    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 deep_usage(self) -> int:

+        return self.usage + sum(c.deep_usage() for c in self.children.values())

-    def leaves(self) -> set["RadixTree"]:

-        if self.children:

-            return {l for c in self.children for l in c.leaves()}

+    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 {self}

+            return None, None

     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

+    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()

-t = RadixTree()

-words = {w.strip() for w in sys.stdin}

+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.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())

+    t.add(w)

+for w in t.traverse():

+    print("".join(w))