fork download
copy 
  1. import sys
  2. from collections import defaultdict
  3. sys.setrecursionlimit(10**7)
  4.  
  5. class SegmentTree:
  6.     def __init__(self, data, k):
  7.         self.n = len(data)
  8.         self.k = k
  9.         self.tree = [0] * (4 * self.n)
  10.         self.build(data, 1, 0, self.n-1)
  11.  
  12.     def build(self, data, node, start, end):
  13.         if start == end:
  14.             self.tree[node] = data[start]
  15.         else:
  16.             mid = (start + end) // 2
  17.             self.build(data, 2*node, start, mid)
  18.             self.build(data, 2*node+1, mid+1, end)
  19.             self.tree[node] = self.tree[2*node] & self.tree[2*node+1]
  20.  
  21.     def query(self, node, start, end, l, r):
  22.         if r < start or end < l:
  23.             return (1 << self.k) - 1
  24.         if l <= start and end <= r:
  25.             return self.tree[node]
  26.         mid = (start + end) // 2
  27.         left = self.query(2*node, start, mid, l, r)
  28.         right = self.query(2*node+1, mid+1, end, l, r)
  29.         return left & right
  30.  
  31. class HLD:
  32.     def __init__(self, n, k, adj, color_sets):
  33.         self.n = n
  34.         self.k = k
  35.         self.adj = adj
  36.         self.color_sets = color_sets
  37.         self.size = [0] * (n+1)
  38.         self.depth = [0] * (n+1)
  39.         self.parent = [0] * (n+1)
  40.         self.heavy = [-1] * (n+1)
  41.         self.head = [0] * (n+1)
  42.         self.pos = [0] * (n+1)
  43.         self.curr_pos = 0
  44.  
  45.         self.dfs(1, 0)
  46.         self.decompose(1, 1)
  47.  
  48.         data = [0] * n
  49.         for u in range(1, n+1):
  50.             bitset = 0
  51.             for c in self.color_sets[u]:
  52.                 bitset |= 1 << (c - 1)
  53.             data[self.pos[u]] = bitset
  54.  
  55.         self.segtree = SegmentTree(data, k)
  56.  
  57.     def dfs(self, u, p):
  58.         self.size[u] = 1
  59.         self.parent[u] = p
  60.         max_size = 0
  61.         for v in self.adj[u]:
  62.             if v == p:
  63.                 continue
  64.             self.depth[v] = self.depth[u] + 1
  65.             self.dfs(v, u)
  66.             self.size[u] += self.size[v]
  67.             if self.size[v] > max_size:
  68.                 max_size = self.size[v]
  69.                 self.heavy[u] = v
  70.  
  71.     def decompose(self, u, h):
  72.         self.head[u] = h
  73.         self.pos[u] = self.curr_pos
  74.         self.curr_pos += 1
  75.         if self.heavy[u] != -1:
  76.             self.decompose(self.heavy[u], h)
  77.         for v in self.adj[u]:
  78.             if v != self.parent[u] and v != self.heavy[u]:
  79.                 self.decompose(v, v)
  80.  
  81.     def query_path(self, u, v):
  82.         res = (1 << self.k) - 1
  83.         while self.head[u] != self.head[v]:
  84.             if self.depth[self.head[u]] < self.depth[self.head[v]]:
  85.                 u, v = v, u
  86.             start = self.pos[self.head[u]]
  87.             end = self.pos[u]
  88.             res &= self.segtree.query(1, 0, self.n - 1, start, end)
  89.             u = self.parent[self.head[u]]
  90.         if self.depth[u] > self.depth[v]:
  91.             u, v = v, u
  92.         res &= self.segtree.query(1, 0, self.n - 1, self.pos[u], self.pos[v])
  93.         return bin(res).count('1')
  94.  
  95. def main():
  96.     t = int(input())
  97.     output = []
  98.     for _ in range(t):
  99.         n, k, s, q = map(int, input().split())
  100.         adj = defaultdict(list)
  101.         for __ in range(n-1):
  102.             u, v = map(int, input().split())
  103.             adj[u].append(v)
  104.             adj[v].append(u)
  105.  
  106.         color_sets = {i: set() for i in range(1, n+1)}
  107.         for __ in range(s):
  108.             v, x = map(int, input().split())
  109.             color_sets[v].add(x)
  110.  
  111.         queries = []
  112.         for __ in range(q):
  113.             u, v = map(int, input().split())
  114.             queries.append((u, v))
  115.  
  116.         hld = HLD(n, k, adj, color_sets)
  117.         results = [str(hld.query_path(u, v)) for u, v in queries]
  118.         output.append(' '.join(results))
  119.  
  120.     print('\n'.join(output))
  121.  
  122. main()
  123.  
Success #stdin #stdout 0.08s 14112KB
comments (?)
stdin
copy 
2
3 5 10 4
1 3
2 1
1 1
1 2
1 3
1 4
1 5
2 1
2 2
2 5
3 1
3 2
1 3
2 3
1 2
1 1
9 3 12 10
7 2
2 4
6 8
9 6
2 1
5 8
2 5
3 9
1 3
6 1
9 3
9 1
5 1
2 3
8 1
4 3
5 3
8 3
7 3
3 1
4 7
1 4
4 5
5 5
4 2
9 9
2 2
2 2
5 2
7 3
stdout
copy 
2 2 3 5
1 1 1 2 1 2 1 1 1 0
https://ideone.com/jBVQYj
language:
Python 3 (python  3.12)
created:
4 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!