Code : Happy Life in University

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

class Tree {
  public:
    int n, timer;
    vector<vector<int>> adj;
    vector<set<int>> children;
    vector<int> tin, tout;
    vector<long long> a, entity;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        tin.resize(n);
        tout.resize(n);
        children.resize(n);
        a.resize(n);
        entity.resize(2 * n + 1);
        timer = 0;
    }

  public:
    void dfs(int src, int par) {
        timer++;
        tin[src] = timer;
        entity[timer] = src;

        for (auto child : adj[src]) {
            if (child != par) {
                children[src].insert(child);
                dfs(child, src);
            }
        }

        timer++;
        tout[src] = timer;
        entity[timer] = src;
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 1; i < n; i++) {
        int pi;
        cin >> pi;
        pi--;
        t.adj[pi].push_back(i);
        t.adj[i].push_back(pi);
    }
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
        t.a[i]--;
    }
    t.dfs(0, -1);

    long long ans = 1;
    for (int i = 0; i < n; i++) {
        // Make sure to exclude tout[i] to avoid empty path.
        // We need to partition the range into disjoint subtrees (per children),
        // take max in each of them and take their best product.
        map<int, int> colors;
        vector<long long> paths = {1};
        int mx = 0;
        for (int x = t.tin[i]; x < t.tout[i]; x++) {
            int entity = t.entity[x];
            int c_now = t.a[entity];

            // It was a discovery event.
            if (t.tin[entity] == x) {
                colors[c_now]++;
                mx = max(mx, (int)(colors.size()));
            } else {
                // It was a finish event.
                colors[c_now]--;
                if (colors[c_now] == 0) {
                    colors.erase(c_now);
                }

                mx = max(mx, (int)(colors.size()));

                // This was a finish event, if it is the finish time of
                // a direct children, store the subtree max.
                if (t.children[i].find(entity) != t.children[i].end()) {
                    paths.push_back(mx);
                    mx = 1;
                }
            }
        }

        sort(paths.begin(), paths.end());
        int sz = paths.size();
        if (paths.size() >= 2) {
            ans = max(ans, paths[sz - 1] * paths[sz - 2]);
        }
    }
    cout << ans << "\n";
}

int main() {
    int t;
    cin >> t;
    for (int zz = 0; zz < t; zz++) {
        solve();
    }
    return 0;
}