Related Articles
Count of distinct colors in a subtree of a Colored Tree with given min frequency for Q queries
• Difficulty Level : Medium
• Last Updated : 04 Jun, 2020

Given a N-ary tree with some colour associated with every node and Q queries. Each query contains two integers A and X. The task is to count all distinct colors in a subtree rooted at A, having frequency of colors greater than or equal to X in that subtree.

Examples:

Input: Tree:


1(1)
/       \
/          \
2(2)         5(3)
/   \       /  |   \
/     \     /   |    \
3(2)    4(3) 6(2) 7(3)  8(3)


query[] = {{1, 2}, {1, 3}, {1, 4}, {2, 3}, {5, 3}}
Output: {2, 2, 1, 0, 1}
Explanation:
In the subtree rooted at 1, the frequency of colour 2 is 3 and colour 3 is 4. So the answer is 2 for the 1 query, 2 for 2nd query and 1 for 3rd query.
For subtree rooted at 2, frequency of color 2 is 2 and color 3 is 1. So no color have frequency more than or equal to 4.
For subtree rooted at 5, frequency of color 2 is 1 and color 3 is 3. So color 3 have frequency equal to 3.

## Recommended: Please try your approach on {IDE} first, before moving on to the solution.

Naive Approach

• For each query, we’ll traverse the whole subtree of the given node.
• We’ll maintain a map which stores the frequency of every colour in the subtree of given node.
• Then, traverse the map and count number of colours such that it’s frequency is greater than given x.

Time Complexity: O(Q * N)
Space Complexity: O(Q * N)

Approach: (Using Mo’s Algorithm)

• We’ll first flat the tree using Euler Tour.
• We’ll give the number to every node, when it will go in the DFS and when it came out. Let’s denote this with tin[node] and tout[node] for every node.
• After we have flatten the tree into an array, every sub tree of can be denoted as some some array with start and end index as tin[node] and tout[node] respectively.
• Now the question changed into number of elements with frequency greater than equal to X in some subarray.
• We’ll use Mo’s algorithm to solve this problem.
• First, we’ll store the queries and sort them according to the tin[node] / SQ where SQ is the Square root of N.
• As we move the pointers, we’ll store the frequency at ith colour in an array and answer to the query is stored at Xth position of array as it stores the count of colours with frequnecy greater equal to X.
 // C++ program to count distinct colors// in a subtree of a Colored Tree// with given min frequency for Q queries  #include using namespace std;  const int N = 1e5 + 5;  // graph as adjacency listvector > v(N);  // colour written on node.vector<int> colour(N);  // order of node entering in DFSvector<int> in(N);  // order of node exiting in DFSvector<int> out(N);  // for counting the frequency of// nodes with colour ivector<int> cnt(N);  // for storing frequency of coloursvector<int> freq(N);  // tree in a flatten// form (using Euler tour)vector<int> flatTree(N);  // number of nodesint n,      // square root of n    sq;  // indexes for in and// out of node in DFSint start = 0;  // DFS function to find// order of euler tourvoid DFSEuler(int a, int par){      // stroing the start index    in[a] = ++start;      // storing colour of node    // in flatten array    flatTree[start] = colour[a];      for (int i : v[a]) {          // doing DFS on its child        // skipping parent        if (i == par)            continue;          DFSEuler(i, a);    }      // out index of the node.    out[a] = start;}  // comparator for queriesbool comparator(pair<int, int>& a,                pair<int, int>& b){    // comparator for queries to be    // sorted accoring to in[x] / sq    if (in[a.first] / sq != in[b.first] / sq)        return in[a.first] < in[b.first];      return out[a.first] < out[b.first];}  // Function to answer the queriesvoid solve(vector > arr,           int q){    sq = sqrt(n) + 1;      // for storing answers    vector<int> answer(q);      // for storing indexes of queries    // in the order of input.    map, int> idx;      for (int i = 0; i < q; i++) {          // storing indexes of queries        idx[arr[i]] = i;    }      // doing depth first search to    // find indexes to flat the    // tree using euler tour.    DFSEuler(1, 0);      // After doing Euler tour,    // subtree of x can be    // represented as a subarray    // from in[x] to out[x];      // we'll sort the queries    // according to the in[i];    sort(arr.begin(),         arr.end(),         comparator);      // two pointers for    // sliding the window    int l = 1, r = 0;      for (int i = 0; i < q; i++) {          // finding answer to the query        int node = arr[i].first,            x = arr[i].second;        int id = idx[arr[i]];          while (l > in[node]) {              // decrementing the pointer as            // it is greater than start            // and adding answer            // to our freq array.            l--;            cnt[flatTree[l]]++;            freq[cnt[flatTree[l]]]++;        }          while (r < out[node]) {              // incrementing pointer as it is            // less than the end value and            // adding answer to our freq array.            r++;            cnt[flatTree[r]]++;            freq[cnt[flatTree[r]]]++;        }          while (l < in[node]) {              // removing the lth node from            // freq array and incrementing            // the pointer            freq[cnt[flatTree[l]]]--;            cnt[flatTree[l]]--;            l++;        }          while (r > out[node]) {              // removing the rth node from            // freq array and decrementing            // the pointer            freq[cnt[flatTree[r]]]--;            cnt[flatTree[r]]--;            r--;        }          // answer to this query        // is stored at freq[x]        // freq[x] stores the frequency        // of nodes greater equal to x        answer[id] = freq[x];    }      // printing the queries    for (int i = 0; i < q; i++)        cout << answer[i] << " ";}  int main(){    // Driver Code    /*               1(1)             /       \            /          \         2(2)         5(3)        /   \       /  |   \       /     \     /   |    \    3(2)    4(3) 6(2) 7(3)  8(3)    */    n = 8;    v.push_back(2);    v.push_back(1);    v.push_back(3);    v.push_back(2);    v.push_back(4);    v.push_back(2);    v.push_back(5);    v.push_back(1);    v.push_back(6);    v.push_back(5);    v.push_back(7);    v.push_back(5);    v.push_back(8);    v.push_back(5);      colour = 1;    colour = 2;    colour = 2;    colour = 3;    colour = 3;    colour = 2;    colour = 3;    colour = 3;      vector > queries        = { { 1, 2 },            { 1, 3 },            { 1, 4 },            { 2, 3 },            { 5, 3 } };    int q = queries.size();      solve(queries, q);    return 0;}
Output:
2 2 1 0 1


Time Complexity: Space Complexity: Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.  Get hold of all the important mathematical concepts for competitive programming with the Essential Maths for CP Course at a student-friendly price.

In case you wish to attend live classes with industry experts, please refer Geeks Classes Live and Geeks Classes Live USA

My Personal Notes arrow_drop_up