Find MEX of every subtree in the given Tree

Introduction

This blog will discuss the solution to the problem to find MEX of every subtree in the given tree. We are given a generic tree that consists of N nodes, and the tree is rooted at node 0, and an array value[] is given, which represents the value of each node. We have to find MEX for a subtree of every node. Before we deep dive into the solution of this problem, we should look at a sample example to better understand this problem.

MEX of a subtree = Smallest non-negative integer which is missing in the subtree, including the node value.

Sample Examples

Input: N = 6, edges = {{0, 1}, {1, 2}, {0, 3}, {3, 4}, {3, 5}}, val[] = {1, 3, 5, 2, 0, 4}

Output:[6, 0, 0, 1, 1, 0] 

Approach

This problem to find mex of every subtree in the given tree can be solved by doing DFS Algorithm traversal on the tree and by also performing the binary search to find the missing minimum positive integer for every node subtree.

1. We will initialize an array solution[] to store the mex of every node subtree.
2. Then we will do dfs traversal from node 0, and we will perform the following steps:

• We will store all the nodes during dfs traversal in the vector.
• We will merge two vectors in the sorted order for each node.
• Now we will do a binary search on the sorted vector to find the mex, and we will store the mex of the current subtree in the solution[] vector.

After performing the above steps, we will print the solution[] vector.  

Implementation in C++

// C++ program to find MEX of every subtree in the given Tree

#include <bits/stdc++.h>
using namespace std;

// to store the edges of the tree
vector<vector<int>> edges;

// to add the edges of the tree
void add_edge(int x, int y)
{
    edges.push_back({x, y});
}

// to merge two sorted arrays
vector<int> merge(vector<int> &a,
                  vector<int> &b)
{
    // pointer for traversing both the arrays
    int i = 0, j = 0;

    // to store the merged array
    vector<int> result;

    int n = a.size(), m = b.size();

    while (i < n && j < m)
    {
        // if a[i]<b[j] then pushback a[i] in the array and increment i pointer
        if (a[i] < b[j])
            result.push_back(a[i++]);
        // else if a[i]>b[j] then pushback b[j] in the array and increment j pointer
        else if (b[j] < a[i])
            result.push_back(b[j++]);
    }

    // if there are elements left in the first array then pushback the remaining elements
    while (i < n)
        result.push_back(a[i++]);

    // if there are elements left in the second array then pushback the remaining elements
    while (j < m)
        result.push_back(b[j++]);

    return result;
}

// to perform the dfs traversal on the tree
vector<int> helper(vector<int> tree[], int x,
                int p, vector<int> &first,
                vector<int> &solution)
{
    vector<int> result;
    result.push_back(first[x]);

    // iterating the childrens
    for (auto i : tree[x])
    {

        // storing all values of the subtree i in sorted order
        if (i != p)
        {
            vector<int> temp = helper(tree, i, x, first, solution);
            result = merge(result, temp);
        }
    }

    int l = 0, r = result.size() - 1;
    int ans = result.size();

    // binary search in the merged array to find the mex
    while (l <= r)
    {
        int mid = (l + r) / 2;

        // update the ranges
        if (result[mid] > mid)
            r = mid - 1;
        else
        {
            ans = mid + 1;
            l = mid + 1;
        }
    }
    if (result[0] != 0)
        ans = 0;

    // update the mex for current tree node
    solution[x] = ans;

    return result;
}

// to find the mex for each subtree of tree
void solve(int N, vector<int> x)
{
    vector<int> tree[N + 1];
    for (auto i : edges)
    {
        tree[i[0]].push_back(i[1]);
        tree[i[1]].push_back(i[0]);
    }
    vector<int> solution(N, 0);

    // Function Call
    helper(tree, 0, -1, x, solution);

    // printing the answer for every node
    for (auto i : solution)
        cout << i << " ";
}

// Driver Code
int main()
{
    int N = 6;
    add_edge(0, 1);
    add_edge(1, 2);
    add_edge(0, 3);
    add_edge(3, 4);
    add_edge(3, 5);

    vector<int> val =  {1, 3, 5, 2, 0, 4};
    solve(N, val);

    return 0;
}

You can also try this code with Online C++ Compiler

Run Code

Output:

6 0 0 1 1 0

Complexity Analysis

Time Complexity: O(N*(N+log(N))

Space Complexity: O(N)