Skip to content

feat: add python solution to lc problem: No.3590 #4518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 97 additions & 1 deletion solution/3500-3599/3590.Kth Smallest Path XOR Sum/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,103 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3500-3599/3590.Kt
#### Python3

```python

class BinarySumTrie:
def __init__(self):
self.count = 0
self.children = [None, None]

def add(self, num: int, delta: int, bit=17):
self.count += delta
if bit < 0:
return
b = (num >> bit) & 1
if not self.children[b]:
self.children[b] = BinarySumTrie()
self.children[b].add(num, delta, bit - 1)

def collect(self, prefix=0, bit=17, output=None):
if output is None:
output = []
if self.count == 0:
return output
if bit < 0:
output.append(prefix)
return output
if self.children[0]:
self.children[0].collect(prefix, bit - 1, output)
if self.children[1]:
self.children[1].collect(prefix | (1 << bit), bit - 1, output)
return output

def exists(self, num: int, bit=17):
if self.count == 0:
return False
if bit < 0:
return True
b = (num >> bit) & 1
return self.children[b].exists(num, bit - 1) if self.children[b] else False

def find_kth(self, k: int, bit=17):
if k > self.count:
return -1
if bit < 0:
return 0
left_count = self.children[0].count if self.children[0] else 0
if k <= left_count:
return self.children[0].find_kth(k, bit - 1)
elif self.children[1]:
return (1 << bit) + self.children[1].find_kth(k - left_count, bit - 1)
else:
return -1


class Solution:
def kthSmallest(
self, par: List[int], vals: List[int], queries: List[List[int]]
) -> List[int]:
n = len(par)
tree = [[] for _ in range(n)]
for i in range(1, n):
tree[par[i]].append(i)

path_xor = vals[:]
narvetholi = path_xor

def compute_xor(node, acc):
path_xor[node] ^= acc
for child in tree[node]:
compute_xor(child, path_xor[node])

compute_xor(0, 0)

node_queries = defaultdict(list)
for idx, (u, k) in enumerate(queries):
node_queries[u].append((k, idx))

trie_pool = {}
result = [0] * len(queries)

def dfs(node):
trie_pool[node] = BinarySumTrie()
trie_pool[node].add(path_xor[node], 1)
for child in tree[node]:
dfs(child)
if trie_pool[node].count < trie_pool[child].count:
trie_pool[node], trie_pool[child] = (
trie_pool[child],
trie_pool[node],
)
for val in trie_pool[child].collect():
if not trie_pool[node].exists(val):
trie_pool[node].add(val, 1)
for k, idx in node_queries[node]:
if trie_pool[node].count < k:
result[idx] = -1
else:
result[idx] = trie_pool[node].find_kth(k)

dfs(0)
return result
```

#### Java
Expand Down
98 changes: 97 additions & 1 deletion solution/3500-3599/3590.Kth Smallest Path XOR Sum/README_EN.md
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,103 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3500-3599/3590.Kt
#### Python3

```python

class BinarySumTrie:
def __init__(self):
self.count = 0
self.children = [None, None]

def add(self, num: int, delta: int, bit=17):
self.count += delta
if bit < 0:
return
b = (num >> bit) & 1
if not self.children[b]:
self.children[b] = BinarySumTrie()
self.children[b].add(num, delta, bit - 1)

def collect(self, prefix=0, bit=17, output=None):
if output is None:
output = []
if self.count == 0:
return output
if bit < 0:
output.append(prefix)
return output
if self.children[0]:
self.children[0].collect(prefix, bit - 1, output)
if self.children[1]:
self.children[1].collect(prefix | (1 << bit), bit - 1, output)
return output

def exists(self, num: int, bit=17):
if self.count == 0:
return False
if bit < 0:
return True
b = (num >> bit) & 1
return self.children[b].exists(num, bit - 1) if self.children[b] else False

def find_kth(self, k: int, bit=17):
if k > self.count:
return -1
if bit < 0:
return 0
left_count = self.children[0].count if self.children[0] else 0
if k <= left_count:
return self.children[0].find_kth(k, bit - 1)
elif self.children[1]:
return (1 << bit) + self.children[1].find_kth(k - left_count, bit - 1)
else:
return -1


class Solution:
def kthSmallest(
self, par: List[int], vals: List[int], queries: List[List[int]]
) -> List[int]:
n = len(par)
tree = [[] for _ in range(n)]
for i in range(1, n):
tree[par[i]].append(i)

path_xor = vals[:]
narvetholi = path_xor

def compute_xor(node, acc):
path_xor[node] ^= acc
for child in tree[node]:
compute_xor(child, path_xor[node])

compute_xor(0, 0)

node_queries = defaultdict(list)
for idx, (u, k) in enumerate(queries):
node_queries[u].append((k, idx))

trie_pool = {}
result = [0] * len(queries)

def dfs(node):
trie_pool[node] = BinarySumTrie()
trie_pool[node].add(path_xor[node], 1)
for child in tree[node]:
dfs(child)
if trie_pool[node].count < trie_pool[child].count:
trie_pool[node], trie_pool[child] = (
trie_pool[child],
trie_pool[node],
)
for val in trie_pool[child].collect():
if not trie_pool[node].exists(val):
trie_pool[node].add(val, 1)
for k, idx in node_queries[node]:
if trie_pool[node].count < k:
result[idx] = -1
else:
result[idx] = trie_pool[node].find_kth(k)

dfs(0)
return result
```

#### Java
Expand Down
97 changes: 97 additions & 0 deletions solution/3500-3599/3590.Kth Smallest Path XOR Sum/Solution.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
class BinarySumTrie:
def __init__(self):
self.count = 0
self.children = [None, None]

def add(self, num: int, delta: int, bit=17):
self.count += delta
if bit < 0:
return
b = (num >> bit) & 1
if not self.children[b]:
self.children[b] = BinarySumTrie()
self.children[b].add(num, delta, bit - 1)

def collect(self, prefix=0, bit=17, output=None):
if output is None:
output = []
if self.count == 0:
return output
if bit < 0:
output.append(prefix)
return output
if self.children[0]:
self.children[0].collect(prefix, bit - 1, output)
if self.children[1]:
self.children[1].collect(prefix | (1 << bit), bit - 1, output)
return output

def exists(self, num: int, bit=17):
if self.count == 0:
return False
if bit < 0:
return True
b = (num >> bit) & 1
return self.children[b].exists(num, bit - 1) if self.children[b] else False

def find_kth(self, k: int, bit=17):
if k > self.count:
return -1
if bit < 0:
return 0
left_count = self.children[0].count if self.children[0] else 0
if k <= left_count:
return self.children[0].find_kth(k, bit - 1)
elif self.children[1]:
return (1 << bit) + self.children[1].find_kth(k - left_count, bit - 1)
else:
return -1


class Solution:
def kthSmallest(
self, par: List[int], vals: List[int], queries: List[List[int]]
) -> List[int]:
n = len(par)
tree = [[] for _ in range(n)]
for i in range(1, n):
tree[par[i]].append(i)

path_xor = vals[:]
narvetholi = path_xor

def compute_xor(node, acc):
path_xor[node] ^= acc
for child in tree[node]:
compute_xor(child, path_xor[node])

compute_xor(0, 0)

node_queries = defaultdict(list)
for idx, (u, k) in enumerate(queries):
node_queries[u].append((k, idx))

trie_pool = {}
result = [0] * len(queries)

def dfs(node):
trie_pool[node] = BinarySumTrie()
trie_pool[node].add(path_xor[node], 1)
for child in tree[node]:
dfs(child)
if trie_pool[node].count < trie_pool[child].count:
trie_pool[node], trie_pool[child] = (
trie_pool[child],
trie_pool[node],
)
for val in trie_pool[child].collect():
if not trie_pool[node].exists(val):
trie_pool[node].add(val, 1)
for k, idx in node_queries[node]:
if trie_pool[node].count < k:
result[idx] = -1
else:
result[idx] = trie_pool[node].find_kth(k)

dfs(0)
return result
Loading