MEX from root node to every other node in a tree

Question

Given a rooted tree having N nodes. Root node is node 1. Each ith node has some value , val[i] associated with it.

For each node i (1<=i<=N) we want to know MEX of the path values from root node to node i.

MEX of an array is smallest positive integer not present in the array, for instance MEX of {1,2,4} is 3

Example : Say we are given tree with 4 nodes. Value of nodes are [1,3,2,8] and we also have parent of each node i (other than node 1 as it is the root node). Parent array is defined as [1,2,2] for this example. It means parent of node 2 is node 1, parent of node 3 is node 2 and parent of node 4 is also node 2.

Node 1 : MEX(1) = 2
Node 2 : MEX(1,3) = 2
Node 3 : MEX(1,3,2) = 4
Node 4 : MEX(1,3,8) = 2

Hence answer is [2,2,4,2]

In worst case total number of Nodes can be upto 10^6 and value of each node can go upto 10^9.

Attempt :

Approach 1 : As we know MEX of N elements will be always be between 1 to N+1. I was trying to use this understanding with this tree problem, but then in this case N will keep on changing dynamically as one proceed towards leaf nodes.

Approach 2 : Another thought was to create an array with N+1 empty values and then try to fill them as we go along from root node. But then challenge I faced was on to keep track of first non filled value in this array.

Answer 1

public class PathMex {
    static void dfs(int node, int mexVal, int[] res, int[] values, ArrayList<ArrayList<Integer>> adj, HashMap<Integer, Integer> map) {
        if (!map.containsKey(values[node])) {
            map.put(values[node], 1);
        }
        else {
            map.put(values[node], map.get(values[node]) + 1);
        }
        while(map.containsKey(mexVal)) mexVal++;
        res[node] = mexVal;

        ArrayList<Integer> children = adj.get(node);
        for (Integer child : children) {
            dfs(child, mexVal, res, values, adj, map);
        }

        if (map.containsKey(values[node])) {
            if (map.get(values[node]) == 1) {
                map.remove(values[node]);
            }
            else {
                map.put(values[node], map.get(values[node]) - 1);
            }
        }
    }

    static int[] findPathMex(int nodes, int[] values, int[] parent) {
        ArrayList<ArrayList<Integer>> adj = new ArrayList<>(nodes);
        HashMap<Integer, Integer> map = new HashMap<>();

        int[] res = new int[nodes];

        for (int i = 0; i < nodes; i++) {
            adj.add(new ArrayList<Integer>());
        }
        for (int i = 0; i < nodes - 1; i++) {
            adj.get(parent[i] - 1).add(i + 1);
        }
        dfs(0, 1, res, values, adj, map);
        return res;
    }

    public static void main(String args[]) {
        Scanner sc = new Scanner(System.in);
        int nodes = sc.nextInt();
        int[] values = new int[nodes];
        int[] parent = new int[nodes - 1];
        for (int i = 0; i < nodes; i++) {
            values[i] = sc.nextInt();
        }
        for (int i = 0; i < nodes - 1; i++) {
            parent[i] = sc.nextInt();
        }
        int[] res = findPathMex(nodes, values, parent);
        for (int i = 0; i < nodes; i++) {
            System.out.print(res[i] + " ");
        }
    }
}

Answer 2

public class TestClass {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        PrintWriter wr = new PrintWriter(System.out);
        int T = Integer.parseInt(br.readLine().trim());
        for(int t_i = 0; t_i < T; t_i++)
        {
            int N = Integer.parseInt(br.readLine().trim());
            String[] arr_val = br.readLine().split(" ");
            int[] val = new int[N];
            for(int i_val = 0; i_val < arr_val.length; i_val++)
            {
                val[i_val] = Integer.parseInt(arr_val[i_val]);
            }
            String[] arr_parent = br.readLine().split(" ");
            int[] parent = new int[N-1];
            for(int i_parent = 0; i_parent < arr_parent.length; i_parent++)
            {
                parent[i_parent] = Integer.parseInt(arr_parent[i_parent]);
            }

            int[] out_ = solve(N, val, parent);
            System.out.print(out_[0]);
            for(int i_out_ = 1; i_out_ < out_.length; i_out_++)
            {
                System.out.print(" " + out_[i_out_]);
            }
            
            System.out.println();
            
         }

         wr.close();
         br.close();
    }
    static int[] solve(int N, int[] val, int[] parent){
       // Write your code here
        int[] result = new int[val.length];
        ArrayList<ArrayList<Integer>> temp = new ArrayList<>();
        ArrayList<Integer> curr = new ArrayList<>();

        if(val[0]==1)
            curr.add(2);
        else{
            curr.add(1);
            curr.add(val[0]);
        }
        result[0]=curr.get(0);
        temp.add(new ArrayList<>(curr));
        for(int i=1;i<val.length;i++){

            int parentIndex = parent[i-1]-1;
            curr = new ArrayList<>(temp.get(parentIndex));
            int nodeValue = val[i];
            boolean enter = false;
            while(curr.size()>0 && nodeValue == curr.get(0)){
                curr.remove(0);
                nodeValue++;
                enter=true;
            }
            if(curr.isEmpty())
                curr.add(nodeValue);
            else if(!curr.isEmpty() && curr.contains(nodeValue) ==false && (enter|| curr.get(0)<nodeValue))
                curr.add(nodeValue);

            Collections.sort(curr);
            temp.add(new ArrayList<>(curr));
            result[i]=curr.get(0);
        }

        return result;
    
    }
}

Answer 3

This can be done in time O(n log n) using augmented BSTs.

Imagine you have a data structure that supports the following operations:

1. insert(x), which adds a copy of the number x.
2. remove(x), which removes a copy of the number x.
3. mex(), which returns the MEX of the collection.

With something like this available, you can easily solve the problem by doing a recursive tree walk, inserting items when you start visiting a node and removing those items when you leave a node. That will make n calls to each of these functions, so the goal will be to minimize their costs.

We can do this using augmented BSTs. For now, imagine that all the numbers in the original tree are distinct; we’ll address the case when there are duplicates later. Start off with Your BST of Choice and augment it by having each node store the number of nodes in its left subtree. This can be done without changing the asymptotic cost of an insertion or deletion (if you haven’t seen this before, check out the order statistic tree data structure). You can then find the MEX as follows. Starting at the root, look at its value and the number of nodes in its left subtree. One of the following will happen:

1. The node’s value k is exactly one plus the number of nodes in the left subtree. That means that all the values 1, 2, 3, …, k are in the tree, so the MEX will be the smallest value missing from the right subtree. Recursively find the MEX of the right subtree. As you do, remember that you’ve already seen the values from 1 to k by subtracting k off of all the values you find there as you encounter them.
2. The node’s value k is at least two more than the number of nodes in the left subtree. That means that the there’s a gap somewhere in the node’s in the left subtree plus the root. Recursively find the MEX of the left subtree.

Once you step off the tree, you can look at the last node where you went right and add one to it to get the MEX. (If you never went right, the MEX is 1).

This is a top-down pass on a balanced tree that does O(1) work per node, so it takes a total of O(log n) work.

The only complication is what happens if a value in the original tree (not the augmented BST) is duplicated on a path. But that’s easy to fix: just add a count field to each BST node tracking how many times it’s there, incrementing it when an insert happens and decrementing it when a remove happens. Then, only remove the node from the BST in the case where the frequency drops to zero.

Overall, each operation on such a tree takes time O(log n), so this gives an O(n log n)-time algorithm for your original problem.