CF1916E Happy Life in University

关于我赛时线段树忘了开四倍空间导致白吃了一发罚时这档逝

原题传送门

约定

$x$ 子树内的叶子称为 $x$ 的叶子。

与 $x$ 颜色相同的点称为 $x$ 的同色点或 $x$ 色点。

所有在 $x$ 子树内的、到 $x$ 的路径上（两端不含）没有 $x$ 的同色点的 $x$ 的同色点称为 $x$ 的极浅同色点。

分析

首先考虑两个点要是没有祖先关系，那必然是选择两个叶子最优。想到枚举 LCA 算贡献，则在每个点 $x$ 得知道 $x$ 的所有叶子到 $x$ 的 diff 值。我们考虑在 dfs 回溯的时候维护这个东西。在我们回溯到一个点 $x$ 的时候，$x$ 的叶子到 $x$ 对应的儿子的 diff 值已经有了。我们接下来要做的就是把 $x$ 的所有 到 $x$ 的路径上没有 $x$ 的同色点的叶子 的 diff 值都加 $1$。为此，我们可以先把 $x$ 的所有叶子的 diff 值都加 $1$，然后枚举 $x$ 的所有极浅同色点，把它们的叶子的 diff 值都减 $1$。于是问题转变成求 $x$ 的所有极浅同色点。

我们反过来考虑。$y$ 是 $x$ 的其中一个极浅同色点，等价于从 $y$ 往上跳，遇到的第一个 $y$ 色点是 $x$。于是我们可以在 dfs 的时候维护一个栈，其中放着所有当前点 $x$ 的祖先。然后我们就只需要在这个栈里找最后一个颜色与 $x$ 相同的点 $y$，那么 $x$ 就一定是 $y$ 的其中一个极浅同色点。对每种颜色开一个栈，这样就可以很方便地找到最后一个当前点的同色点。

求出每个点 $x$ 到其所有叶子的 diff 值之后，我们枚举 $x$ 的所有儿子 $y$，对于每个 $y$ 找到 $x$ 到其叶子的 diff 的最大值，然后把这些最大值中的最大值和次大值相乘就是 $x$ 作为 LCA 的最大贡献。

然后我们考虑两个点如果有祖先关系，那必然是一个是根，一个是叶子。直接找根到所有叶子的 diff 最大值即可。

最后还剩下把 $x$ 的所有叶子的 diff 加 / 减 $1$，求子树内叶子 diff 值 max。直接对叶子的 dfs 序开一个线段树，支持区间加区间 max 就没了。

至于为什么对每个点枚举所有它的极浅同色点不会 T，那是因为每个点至多是一个点的极浅同色点。所以每个点最多被枚举一次。

总复杂度是 $O(n\log n)$。

代码

#include <iostream>
#include <cassert>
#include <vector>
#define ll long long
using namespace std;
vector<int> vec[300005];
vector<int> sm[300005];
int clr[300005];
int f[300005];
int head[300005], nxt[300005], to[300005], ecnt;
void add(int u, int v) { to[++ecnt] = v, nxt[ecnt] = head[u], head[u] = ecnt; }
int L[300005], R[300005], ncnt;
ll ans = 1;
int lcnt = 0;
int n;
struct Segment_Tree {
    int mx[1200005], tg[1200005];
    inline void tag(int o, int t) { mx[o] += t, tg[o] += t; }
    inline void pushdown(int o) {
        if (!tg[o]) 
            return;
        tag(o << 1, tg[o]);
        tag(o << 1 | 1, tg[o]);
        tg[o] = 0;
    }
    inline void pushup(int o) { mx[o] = max(mx[o << 1], mx[o << 1 | 1]); }
    void Clear(int o, int l, int r) {
        mx[o] = tg[o] = 0;
        if (l == r) 
            return;
        int mid = (l + r) >> 1;
        Clear(o << 1, l, mid);
        Clear(o << 1 | 1, mid + 1, r);
    }
    void Add(int o, int l, int r, int L, int R, int v) {
        if (L <= l && r <= R) {
            tag(o, v);
            return;
        }
        pushdown(o);
        int mid = (l + r) >> 1;
        if (L <= mid) 
            Add(o << 1, l, mid, L, R, v);
        if (R > mid) 
            Add(o << 1 | 1, mid + 1, r, L, R, v);
        pushup(o);
    }
    int Query(int o, int l, int r, int L, int R) {
        if (L <= l && r <= R) 
            return mx[o];
        pushdown(o);
        int mid = (l + r) >> 1;
        if (R <= mid) 
            return Query(o << 1, l, mid, L, R);
        else if (L > mid) 
            return Query(o << 1 | 1, mid + 1, r, L, R);
        else 
            return max(Query(o << 1, l, mid, L, R), Query(o << 1 | 1, mid + 1, r, L, R));
    }
} seg;
void dfs(int x) {
    if (!vec[clr[x]].empty()) {
        int t = vec[clr[x]].size() * 1 - 1;
        sm[vec[clr[x]][t]].push_back(x);
    }
    if (!head[x]) {
        L[x] = R[x] = ++ncnt;
        seg.Add(1, 1, lcnt, ncnt, ncnt, 1);
        return;
    }
    vec[clr[x]].push_back(x);
    for (int i = head[x]; i; i = nxt[i]) {
        int v = to[i];
        dfs(v);
        R[x] = R[v];
        L[x] = (i == head[x] ? L[v] : L[x]);
    }
    vec[clr[x]].pop_back();
    seg.Add(1, 1, lcnt, L[x], R[x], 1);
    for (auto v : sm[x]) seg.Add(1, 1, lcnt, L[v], R[v], -1);
    int mx = 1, smx = 0;
    for (int i = head[x]; i; i = nxt[i]) {
        int v = to[i];
        int t = seg.Query(1, 1, lcnt, L[v], R[v]);
        if (t >= mx)
            smx = mx, mx = t;
        else if (t > smx) 
            smx = t;
    }
    ans = max(ans, 1ll * mx * smx);
}
signed main() {
    int tc;
    cin >> tc;
    while (tc--) {
        cin >> n;
        for (int i = 0; i <= n + 1; i++) {
            head[i] = 0;
            vec[i].clear();
            sm[i].clear();
            L[i] = R[i] = 0;
        }
        ans = 1;
        lcnt = ncnt = ecnt = 0;
         for (int i = 2; i <= n; i++) {
            cin >> f[i];
            add(f[i], i);
        }
        for (int i = 1; i <= n; i++) lcnt += (head[i] == 0);
        for (int i = 1; i <= n; i++) cin >> clr[i];
        dfs(1);
        cout << ans << "\n";
        seg.Clear(1, 1, lcnt);
    }
    return 0;
}