Chef and Pairs Codechef Solution

by

Chef and Pairs Codechef Solution

You are given a tree (connected, undirected, acyclic graph) consisting of NN nodes. Based on this tree, you have to answer QQ queries.

Each query is of the form:

  • K D V1 V2 ⋯ VKK D V1 V2 ⋯ VK – output the number of pairs (i,j)(i,j), 1≤i<j≤K1≤i<j≤K, such that the shortest path between nodes ViVi and VjVj in the tree has DD edges.

See Also : July Long Challenge 2021 Solutions Codechef

Input

  • The first line contains an integer TT, the number of test cases. Then the test cases follow.
  • The first line of each test case contains two integers, NN and QQ.
  • N−1N−1 lines follow. Each line consists of two integers uu and vv, indicating that there is an edge between nodes uu and vv in the tree.
  • QQ lines follow. Each line describes a query in the format given above.

Output

For each query, output the answer on a new line.

Constraints

  • 1≤T≤51≤T≤5
  • 1≤N,Q≤1051≤N,Q≤105
  • 1≤u,v,Vi≤N1≤u,v,Vi≤N
  • 0≤D≤1050≤D≤105
  • The sum of KK across all queries in a single test case is at most 105105.

Subtasks

  • Subtask 1 (20 points): For each query, K≤10K≤10
  • Subtask 2 (80 points): Original constraints

Sample Input

1
5 2
1 2
1 3
2 4
4 5
3 2 2 3 5
2 4 1 3

Sample Output

2
0

Explanation

In the first query, the pairs of nodes (2,3)(2,3) and (2,5)(2,5) have distance 22.

In the second query, there is no pair with distance 44.

Solution


import math  
def bfs(tree,start,level,parent):
   dp = {}
   par = -1 
   queue = []
   queue.append(start)
   parent[start][0] = -1 
   level[start] = 1 
   while queue:
      temp = queue.pop(0)
      par = parent[temp][0]
      for child in tree[temp]:
         if child!=par:
            queue.append(child)
            parent[child][0] = temp
            level[child] = level[temp]+1
   
   for j in range(1,len(parent[0])):
      for node in range(1,len(level)):
         par = parent[node][j-1]
         if par!=-1:
            parent[node][j] = parent[par][j-1]

def distance(tree,parent,level,a,b,maxsize):
   node1,node2 = a,b 
   if level[a]>level[b]:
      a,b = b,a 
   d = level[b]-level[a] 
   while d!=0:
      i = math.floor(math.log(d)/math.log(2))
      b = parent[b][i]
      d = d - (1<<i)
   if a==b:
      return (level[node1]-level[a])+(level[node2]-level[a])

   for i in range(maxsize-1,-1,-1):
      if (parent[a][i]!=-1) and (parent[a][i] != parent[b][i]):
         a,b = parent[a][i],parent[b][i]
   par = parent[a][0]
   return (level[node1]-level[par])+(level[node2]-level[par])
            


for _ in range(int(input())):
   nodes,queries = map(int,input().split())
   tree = {}
   maxsize = 1+math.ceil(math.log(nodes)/math.log(2))
   parent = [[-1 for _ in range(maxsize)] for _ in range(nodes+1)]
   level = [0]*(nodes+1)
   for i in range(1,nodes+1):
      tree[i] = []
   for _ in range(nodes-1):
      u,v = map(int,input().split())
      tree[u].append(v)
      tree[v].append(u)
   
   bfs(tree,1,level,parent)

   
   for _ in range(queries):
      arr = list(map(int,input().split()))
      k = arr[0]
      d = arr[1]
      arr.pop(0)
      arr.pop(0)
      sm = 0 
      
      for i in range(k):
         for j in range(i+1,k):
            if distance(tree,parent,level,arr[i],arr[j],maxsize)==d:
               sm+=1
      print(sm)

July Long Challenge 2021 Solutions

Weekly Contest 247

Biweekly Contest 55

Leave a Comment

six + thirteen =