Last Updated: 20 Nov, 2021

Catch Pokemon

Ninja
Asked in company
Ola

Problem statement

There are ‘N’ cities in Pokeland, and they are connected by ‘N-1’ roads. It is possible to reach from any city to any other city and there is a unique path between each pair of cities.

Each city has one Pokemon represented by an integer between 1 and 10^9. In each city. Given queries in the form of pairs of cities ‘U’ and ‘V’, find the number of Pokemon Ash catches while traveling from city ‘U’ to ‘V’. If he already has a Pokemon ‘X’ and he comes across it again, he does NOT catch it.

Note: Each query is independent, that is, he starts each query with zero Pokemon.

Input Format:
The first line contains 'T', denoting the number of test cases.
For each Test :
    The first line contains two space separated integers 'N' and ‘Q’, denoting the number of cities and the number of queries respectively.
The next N-1 lines contain a pair of cities ‘X’ and ‘Y’, indicating that there is a road between ‘X’ and ‘Y’.
The next line contains an array ‘P’ of ‘N’ integers, representing the type of Pokemon found in each city. 
The next ‘Q’ lines contain a pair of cities ‘U’ and ‘V’, indicating each query.
Output Format:
Print Q lines, i.e, for each query print the number of Pokemon Ash catches in that travel.
Note:
You are not required to print the expected output. It has already been taken care of. Just implement the function.
Constraints -
1 <= 'T' <= 10
2 <= 'N' <= 10^4
1 <= ‘Q’ <= 10^5
1 <= X,Y,U,V <= N
1<= A[i] <= 10^9, i ∈ (1,N) 

Note - The sum of 'N' over all test cases does not exceed 10^4 and the sum of ‘Q’ does not exceed 10^5.

Time Limit: 1 sec

Approaches

01 Approach

Approach: For each query, we simulate the entire path between both the cities and count the number of unique Pokemon encountered on the path. This is our answer.

To efficiently be able to find the path between any two cities, say ‘A’ and ‘B’, we need to find their Lowest Common Ancestor (LCA). We can compute the LCA in multiple ways, including binary lifting and segment trees.

Once the LCA is computed, we can divide the journey into two paths, namely A to LCA and LCA to B (Corner case is when A or B itself is the LCA). We can traverse both these paths while using a set or a hashmap to keep track of the number of unique values.


 

Algorithm:

  • Apply a DFS on the tree and precompute the Binary lifting array as follows:
    • up[i][j] = up[up[i-1][j]][j-1]
  • Now, to find LCA for any pair of cities, we can use the Binary lifting table to find the Lowest node which is the ancestor of both the nodes.
  • For each pair of cities ‘A’, ‘B’, assign L = lca(A,B)
  • Assign Cur = A
  • While Cur is not equal to L
    • Insert Pokemon at Cur into a set
  • Assign Cur = B
  • While Cur is not equal to L
    • Insert Pokemon at Cur into a set
  • Assign Pokemon at L into the set
  • As a Set does not contain duplicate values, the size of the set is our answer.

02 Approach

Approach: Here, we use Mo’s algorithm to optimize the queries. Mo’s algorithm reorders the queries in the following way:

Splits the array into blocks of size √N and sorts the queries according to which block the start point lies in. If both start points belong in the same block, it sorts them according to their endpoints.

We can flatten the tree into an array using DFS such that each query becomes a query on a range in the array instead of a path on the tree. 

We can now maintain a frequency array or a map to count the occurrences of each value in our range.


 

Algorithm:

  • Apply a DFS on the tree and precompute the Binary lifting array as follows:
    • up[i][j] = up[up[i-1][j]][j-1]
  • While doing the same DFS, also compute the Euler tour of the array, with the In-time and Out-time of each node in the array.
  • Now, to find LCA for any pair of cities, we can use the Binary lifting table to find the Lowest node which is the ancestor of both the nodes.
  • For each pair of cities ‘A’, ‘B’, assign L = lca(A,B)
  • If ‘A’ occurs after ‘B’ in the Euler tour, swap A and B.
  • If A is equal to L
    • The range of path would be [ Intime[A], Intime[B] ]
  • Else
    • The range of path would be [ Outtime[A], Intime[B] ] + [Intime[L], Intime[L]]
  • Now that we have all the ranges, we can sort them in the based on Mo’s algorithm requirement given above and apply Mo’s Algorithm
  • Note, in the given ranges, we ignore the nodes occurring twice.