Separating Numbers

We are given a tree with $N$ nodes. Each node has a color $C_i$. We are also given $N-1$ queries and in each query we are told to destroy a previously undestroyed edge. Every time we destroy an edge, we have to report the number of pairs of nodes that get disconnected. Here, two nodes $i$ and $j$ are said to be disconnected  if before the destruction you could reach from $i$ to $j$ using edges not yet destroyed , and if $C_i=C_j$.

Constraint:

$N\le 300,000$

$C[i]\le 100,000$

Example:

Input:
5
5 4
2 5
3 2
1 2
5 1 5 4 5
2 1 3 4
Output:
2
0
1
0

Approach

Java

import java.util.Map;
import java.util.TreeMap;

public class SeparatingNumbers {
    private static int C = 312345;
    private static long[] ans = new long[C];

    public static void main(String[] args) {
        int N = 5;

        int ar2D[][] = { { 5, 4 }, { 2, 5 }, { 3, 2 }, { 1, 2 } };
        int[] u = new int[C];
        int[] v = new int[C];
        for (int i = 0; i < N - 1; i++) {
            u[i] = ar2D[i][0] - 1;
            v[i] = ar2D[i][1] - 1;
        }
        int[] c = { 5, 1, 5, 4, 5 };
        // in order -1
        int[] order = { 1, 0, 2, 3 };
        DisjointSetUnion dsu = new DisjointSetUnion(N, c);
        for (int i = N - 2; i >= 0; i--) {
            int x = u[order[i]];
            int y = v[order[i]];
            ans[i] = dsu.merge(x, y);
        }
        for (int i = 0; i < N - 1; i++) {
            System.out.println(ans[i]);
        }
    }

    static public class DisjointSetUnion {

        private int rank[], parent[], size[];
        private TreeMap<Integer, Integer> map[];
        private int n;

        public DisjointSetUnion(int n, int c[]) {
            this.n = n;
            makeSet(c);
        }

        private void makeSet(int c[]) {
            rank = new int[n];
            parent = new int[n];
            map = new TreeMap[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                size[i] = 1;
                map[i] = new TreeMap<>();
                increment(map[i], c[i], 1);
            }
        }

        public int find(int x) {
            if (parent[x] != x)
                parent[x] = find(parent[x]);
            return parent[x];
        }

        public long merge(int x, int y) {
            int xRoot = find(x);
            int yRoot = find(y);

            if (xRoot == yRoot)
                return getCount(map[xRoot], map[xRoot]);

            size[xRoot] = size[yRoot] = size[xRoot] + size[yRoot];

            long count = getCount(map[xRoot], map[yRoot]);

            if (rank[xRoot] < rank[yRoot]) {
                parent[xRoot] = yRoot;
                merge(map[yRoot], map[xRoot]);
            } else {
                parent[yRoot] = xRoot;
                merge(map[xRoot], map[yRoot]);
                if (rank[xRoot] == rank[yRoot]) {
                    rank[xRoot]++;
                }
            }

            return count;
        }

        // into x
        private void merge(TreeMap<Integer, Integer> x, TreeMap<Integer, Integer> y) {
            for (Map.Entry<Integer, Integer> entry : y.entrySet()) {
                increment(x, entry.getKey(), entry.getValue());
            }
        }

        private long getCount(TreeMap<Integer, Integer> x, TreeMap<Integer, Integer> y) {
            if (x.size() > y.size()) {
                return getCount(y, x);
            }
            long ans = 0;
            for (Map.Entry<Integer, Integer> entry : x.entrySet()) {
                ans += (long) entry.getValue() * get(y, entry.getKey());
            }
            return ans;
        }
    }

    static void increment(TreeMap<Integer, Integer> map, int key, int x) {
        Integer value = map.get(key);
        if (value == null) {
            map.put(key, x);
        } else {
            map.put(key, value + x);
        }
    }

    static int get(TreeMap<Integer, Integer> map, int key) {
        Integer value = map.get(key);
        if (value == null) {
            return 0;
        }
        return value;
    }

}