Count of distinct colors in a subtree of a Colored Tree with given min frequency for Q queries

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.

Naive Approach



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

Approach: (Using Mo’s Algorithm)

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ program to count distinct colors
// in a subtree of a Colored Tree
// with given min frequency for Q queries
  
#include <bits/stdc++.h>
using namespace std;
  
const int N = 1e5 + 5;
  
// graph as adjacency list
vector<vector<int> > v(N);
  
// colour written on node.
vector<int> colour(N);
  
// order of node entering in DFS
vector<int> in(N);
  
// order of node exiting in DFS
vector<int> out(N);
  
// for counting the frequency of
// nodes with colour i
vector<int> cnt(N);
  
// for storing frequency of colours
vector<int> freq(N);
  
// tree in a flatten
// form (using Euler tour)
vector<int> flatTree(N);
  
// number of nodes
int n,
  
    // square root of n
    sq;
  
// indexes for in and
// out of node in DFS
int start = 0;
  
// DFS function to find
// order of euler tour
void 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 queries
bool 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 queries
void solve(vector<pair<int, int> > 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<pair<int, int>, 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[1].push_back(2);
    v[2].push_back(1);
    v[2].push_back(3);
    v[3].push_back(2);
    v[2].push_back(4);
    v[4].push_back(2);
    v[1].push_back(5);
    v[5].push_back(1);
    v[5].push_back(6);
    v[6].push_back(5);
    v[5].push_back(7);
    v[7].push_back(5);
    v[5].push_back(8);
    v[8].push_back(5);
  
    colour[1] = 1;
    colour[2] = 2;
    colour[3] = 2;
    colour[4] = 3;
    colour[5] = 3;
    colour[6] = 2;
    colour[7] = 3;
    colour[8] = 3;
  
    vector<pair<int, int> > queries
        = { { 1, 2 },
            { 1, 3 },
            { 1, 4 },
            { 2, 3 },
            { 5, 3 } };
    int q = queries.size();
  
    solve(queries, q);
    return 0;
}
chevron_right

Output:
2 2 1 0 1

Time Complexity:
Space Complexity:





Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.


Article Tags :