Refine trie.py by advice of GPT (d6c605c)

trie.py

@@ -1,19 +1,25 @@
 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}
+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()
+        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.height = 1
         self.usage = 0
 
     def __str__(self):
@@ -25,39 +31,28 @@ class RadixTree:
     def add(self, word: str):
         common = os.path.commonprefix([self.label, word])
         suffix = word[len(common):]
-        match common:
-            case self.label:
-                if self.children:
+        if common == self.label:
             for child in self.children:
                 if child.label:
-                            ah = child.add(suffix)
-                            if ah:
-                                nh = max(self.height, 1 + ah)
+                    if child.add(suffix):
                         break
             else:
+                if not self.children:
+                    self.children.add(RadixTree("", self))
                 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:
+        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
-            self.height = nh
-            return nh
+            return True
 
     def full(self) -> str:
         s = self.label


...	...	@@ -1,19 +1,25 @@
1	1	import os
2	2	import sys
3	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}
	4	+def minima(options, key):
	5	+ min_val = None
	6	+ result = set()
	7	+ for o in options:
	8	+ k = key(o)
	9	+ if min_val is None or k < min_val:
	10	+ min_val = k
	11	+ result = {o}
	12	+ elif k == min_val:
	13	+ result.add(o)
	14	+ return result
8	15
9	16	class RadixTree:
10	17	def __init__(self, label: str = None, parent: "RadixTree" = None):
11	18	self.label = label or ""
12		- self.children = set()
	19	+ self.children: set[RadixTree] = set()
13	20	self.parent = parent
14	21	self.weight = 1 if label is not None else 0
15	22	self.depth = parent.depth + 1 if parent else 0
16		- self.height = 1
17	23	self.usage = 0
18	24
19	25	def __str__(self):
...	...	@@ -25,39 +31,28 @@ class RadixTree:
25	31	def add(self, word: str):
26	32	common = os.path.commonprefix([self.label, word])
27	33	suffix = word[len(common):]
28		- match common:
29		- case self.label:
30		- if self.children:
	34	+ if common == self.label:
31	35	for child in self.children:
32	36	if child.label:
33		- ah = child.add(suffix)
34		- if ah:
35		- nh = max(self.height, 1 + ah)
	37	+ if child.add(suffix):
36	38	break
37	39	else:
	40	+ if not self.children:
	41	+ self.children.add(RadixTree("", self))
38	42	self.children.add(RadixTree(suffix, self))
39		- nh = self.height
40		- else:
41		- self.children = {RadixTree(suffix, self), RadixTree("", self)}
42		- nh = 2
43		- case "":
44		- nh = None
45		- case c:
46		- rem = self.label[len(common):]
47		- nc = RadixTree(rem, self)
48		- nc.weight = self.weight
49		- nc.height = self.height
50		- nc.usage = self.usage
51		- nc.children = self.children
52		- for child in nc.children:
53		- child.parent = nc
54		- self.label = c
55		- self.children = {RadixTree(suffix, self), nc}
56		- nh = 1 + self.height
57		- if nh:
	43	+ elif common:
	44	+ remainder = self.label[len(common):]
	45	+ split = RadixTree(remainder, self)
	46	+ split.weight = self.weight
	47	+ split.usage = self.usage
	48	+ split.children = self.children
	49	+ for child in split.children:
	50	+ child.parent = split
	51	+ self.label = common
	52	+ self.children = {RadixTree(suffix, self), split}
	53	+ if common or not self.label:
58	54	self.weight += 1
59		- self.height = nh
60		- return nh
	55	+ return True
61	56
62	57	def full(self) -> str:
63	58	s = self.label
64	59

dost