dost

Implement trie-based method that may be very fast
Parent 8495c0a
dkl9

dkl9 commited on 2025-195 22:52:18
Showing 1 changed files, with 106 additions and 0 deletions.

trie.py
View file @ 0f4f8f8
... ... 
@@ -0,0 +1,106 @@
1 
+import os
2 
+import sys
3 
+
4 
+def minima(options, key) -> set:
5 
+    xy = {o: key(o) for o in options}
6 
+    my = min(xy.values())
7 
+    return {x for (x, y) in xy.items() if y == my}
8 
+
9 
+class RadixTree:
10 
+    def __init__(self, label: str = None, parent: "RadixTree" = None):
11 
+        self.label = label or ""
12 
+        self.children = set()
13 
+        self.parent = parent
14 
+        self.weight = 1 if label is not None else 0
15 
+        self.height = 1
16 
+        self.usage = 0
17 
+
18 
+    def __str__(self):
19 
+        s = f"{self.label or "∅"}"
20 
+        if self.children:
21 
+            s = f"({s} {" ".join(str(x) for x in self.children)})"
22 
+        return s
23 
+
24 
+    def add(self, word: str):
25 
+        common = os.path.commonprefix([self.label, word])
26 
+        suffix = word[len(common):]
27 
+        match common:
28 
+            case self.label:
29 
+                if self.children:
30 
+                    for child in self.children:
31 
+                        if child.label:
32 
+                            ah = child.add(suffix)
33 
+                            if ah:
34 
+                                nh = max(self.height, 1 + ah)
35 
+                                break
36 
+                    else:
37 
+                        self.children.add(RadixTree(suffix, self))
38 
+                        nh = self.height
39 
+                else:
40 
+                    self.children = {RadixTree(suffix, self), RadixTree("", self)}
41 
+                    nh = 2
42 
+            case "":
43 
+                nh = None
44 
+            case c:
45 
+                rem = self.label[len(common):]
46 
+                nc = RadixTree(rem, self)
47 
+                nc.weight = self.weight
48 
+                nc.height = self.height
49 
+                nc.usage = self.usage
50 
+                nc.children = self.children
51 
+                for child in nc.children:
52 
+                    child.parent = nc
53 
+                self.label = c
54 
+                self.children = {RadixTree(suffix, self), nc}
55 
+                nh = 1 + self.height
56 
+        if nh:
57 
+            self.weight += 1
58 
+            self.height = nh
59 
+            return nh
60 
+
61 
+    def full(self) -> str:
62 
+        s = self.label
63 
+        if self.parent:
64 
+            s = self.parent.full() + s
65 
+        return s
66 
+
67 
+    def del_empty(self):
68 
+        for child in list(self.children):
69 
+            if not child.label:
70 
+                if child.children:
71 
+                    child.del_empty()
72 
+                else:
73 
+                    self.children.remove(child)
74 
+
75 
+    def leaves(self) -> set["RadixTree"]:
76 
+        if self.children:
77 
+            return {l for c in self.children for l in c.leaves()}
78 
+        else:
79 
+            return {self}
80 
+
81 
+    def use(self):
82 
+        assert self.usage < self.weight
83 
+        self.usage += 1
84 
+        if self.parent:
85 
+            self.parent.use()
86 
+
87 
+    def best(self) -> "RadixTree":
88 
+        if self.children:
89 
+            branches = minima(self.children, lambda c: (c.usage / c.weight, -c.weight))
90 
+            return min((b.best() for b in branches), key=lambda b: (len(b.full()), b.full()))
91 
+        else:
92 
+            return self
93 
+
94 
+t = RadixTree()
95 
+words = {w.strip() for w in sys.stdin}
96 
+for w in words:
97 
+    t.add(w.strip())
98 
+    t.del_empty()
99 
+leaves = {l.full() for l in t.leaves()}
100 
+assert leaves <= words
101 
+assert words <= leaves
102 
+while t.usage < t.weight:
103 
+    b = t.best()
104 
+    print(b.full())
105 
+    b.use()
106 
+print(str(t))
0 107