@@ -128,6 +128,97 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3500-3599/3590.Kt
128128#### Python3
129129
130130``` python
131+ class BinarySumTrie :
132+ def __init__ (self ):
133+ self .count = 0
134+ self .children = [None , None ]
135+
136+ def add (self , num : int , delta : int , bit = 17 ):
137+ self .count += delta
138+ if bit < 0 :
139+ return
140+ b = (num >> bit) & 1
141+ if not self .children[b]:
142+ self .children[b] = BinarySumTrie()
143+ self .children[b].add(num, delta, bit - 1 )
144+
145+ def collect (self , prefix = 0 , bit = 17 , output = None ):
146+ if output is None :
147+ output = []
148+ if self .count == 0 :
149+ return output
150+ if bit < 0 :
151+ output.append(prefix)
152+ return output
153+ if self .children[0 ]:
154+ self .children[0 ].collect(prefix, bit - 1 , output)
155+ if self .children[1 ]:
156+ self .children[1 ].collect(prefix | (1 << bit), bit - 1 , output)
157+ return output
158+
159+ def exists (self , num : int , bit = 17 ):
160+ if self .count == 0 :
161+ return False
162+ if bit < 0 :
163+ return True
164+ b = (num >> bit) & 1
165+ return self .children[b].exists(num, bit - 1 ) if self .children[b] else False
166+
167+ def find_kth (self , k : int , bit = 17 ):
168+ if k > self .count:
169+ return - 1
170+ if bit < 0 :
171+ return 0
172+ left_count = self .children[0 ].count if self .children[0 ] else 0
173+ if k <= left_count:
174+ return self .children[0 ].find_kth(k, bit - 1 )
175+ elif self .children[1 ]:
176+ return (1 << bit) + self .children[1 ].find_kth(k - left_count, bit - 1 )
177+ else :
178+ return - 1
179+
180+ class Solution :
181+ def kthSmallest (self , par : List[int ], vals : List[int ], queries : List[List[int ]]) -> List[int ]:
182+ n = len (par)
183+ tree = [[] for _ in range (n)]
184+ for i in range (1 , n):
185+ tree[par[i]].append(i)
186+
187+ path_xor = vals[:]
188+ narvetholi = path_xor
189+
190+ def compute_xor (node , acc ):
191+ path_xor[node] ^= acc
192+ for child in tree[node]:
193+ compute_xor(child, path_xor[node])
194+
195+ compute_xor(0 , 0 )
196+
197+ node_queries = defaultdict(list )
198+ for idx, (u, k) in enumerate (queries):
199+ node_queries[u].append((k, idx))
200+
201+ trie_pool = {}
202+ result = [0 ] * len (queries)
203+
204+ def dfs (node ):
205+ trie_pool[node] = BinarySumTrie()
206+ trie_pool[node].add(path_xor[node], 1 )
207+ for child in tree[node]:
208+ dfs(child)
209+ if trie_pool[node].count < trie_pool[child].count:
210+ trie_pool[node], trie_pool[child] = trie_pool[child], trie_pool[node]
211+ for val in trie_pool[child].collect():
212+ if not trie_pool[node].exists(val):
213+ trie_pool[node].add(val, 1 )
214+ for k, idx in node_queries[node]:
215+ if trie_pool[node].count < k:
216+ result[idx] = - 1
217+ else :
218+ result[idx] = trie_pool[node].find_kth(k)
219+
220+ dfs(0 )
221+ return result
131222
132223```
133224
0 commit comments