dost

@@ -0,0 +1,2 @@

+*.txt

+tags

@@ -1,6 +1,7 @@

 from collections import defaultdict

 import re

 import sys

+import time

 from typing import Generator

 def minima(options, key):

@@ -15,15 +16,18 @@ def minima(options, key):

             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

+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}]"

@@ -32,25 +36,17 @@ class Trie:

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

-            )

+            token, child = min(self.children.items(), key=cost)

             suffix, leaf = child.best()

             if not leaf:

                 return None, None

@@ -58,7 +54,11 @@ class Trie:

         else:

             return None, None

-    def use(self):

+    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

@@ -66,16 +66,16 @@ class Trie:

         while True:

             tokens, leaf = self.best()

             if not leaf:

-                assert self.deep_usage() == self.deep_weight()

+                assert self.deep_usage == self.deep_weight

                 break

             yield tokens

-            leaf.use()

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

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