PROBLEM LINK:
Practice
Contest
Author: GVS Akhil
Testers: GVS Akhil, Kumar Abhinav
Editorialist: GVS Akhil
DIFFICULTY:
Medium
PREREQUISUITES:
Disjoint Set Union(DSU) on Tree (Refer here
)
PROBLEM
Given queries specifying node X and colour Z , find the number of ordered triplets of nodes (A ,B, C) where A is an ancestor of X , B & C are nodes in the subtree of X satisfying LCA(B,C)=X and colour_A = colour_B = colour_C = Z.
(A,B,C) are also pairwise distinct.
EXPLANATION
Since there are no persistent updates, we can attempt to solve this problem offline using DSU on Tree method mentioned in the link in prerequisites.
We can maintain a map at each node during processing which stores the number of nodes having some particular colour in its subtree.
During the map merge operation when performing DSU, we would need to keep track of the number of valid ordered pairs of (B,C) that we can pick for any colour. We can calculate this in 2 ways:
Approach 1:
Let us fix node X and colour Z. Let the children of X be C_1,C_2...C_k. Let the count of nodes with colour Z in C_i be P_i.
The contribution from each C_i to the number of ordered pairs (B,C) are P_i*(\sum_{j=1,j \neq i}^{k} P_j) .
Knowing the above formula, for a fixed X and Z , we can calculate the ordered pairs in O(K*logN) per query (log factor comes from map access) by iterating over each child each time but this is too slow since k can can be equal to N but we can speed it up.
Observation: \sum_{i=1}^{k} (P_i*( \sum_{j=1,j \neq i}^{k} P_j)) = (\sum_{i=1}^k P_i)^2 - \sum_{i=1}^{k} P_i^2
Now we can just maintain another map which keeps track of the sum of square counts and calculate ordered pairs in O(logN) per query.
Approach 2:
This time, let us calculate unordered pairs instead.
Again, Let us fix node X and colour Z. Let the children of X be C_1,C_2...C_k. Let the count of nodes with colour Z in C_i be P_i.
Everytime we encounter C_i, we can pair up its P_i nodes with any of the previously encountered nodes of same colour. That is, for C_i, its contribution to the unordered pairs are P_i *( \sum_{j=1}^{i-1} P_j) . We can easily maintain a running sum to get rid of the summation.
Once we get the number of unordered pairs, we can multiply this number by 2 to get the total ordered pairs.
We can again achieve O(logN) per query using this approach.
Once we obtain the total number of ordered pair for each query using either of the two approaches, we need to multiply this value by the number of ancestors A which have the same colour Z to get the total number of satisfying ordered triplets (A,B,C).
We can maintain a global frequency map to calculate this value. Every time we enter a node, we increment the frequency counter of colour_{node} by 1 and every time we leave a node, we decrement the frequency counter of colour_{node} by 1.
Thus, at any node X, the frequency counter only maintains the counts of colours of ancestors of X and we can find the number of ancestors with colour Z.
This step is again O(logN) per query due to map access.
Final Complexity: O(Nlog^2N + QlogN)
Additional Note: Complexity can be reduced to O(NlogN + Q) by compressing the colour values, using a vector/HLD style DSU and replacing frequency map with frequency array.