@@ -128,6 +128,97 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3500-3599/3590.Kt
128
128
#### Python3
129
129
130
130
``` 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
131
222
132
223
```
133
224
0 commit comments