E. 快乐的大学生活 题解

问题描述

给定一棵以 $1$ 为根的树，共 $n$ 个节点。每个节点 $i$ 有一个颜色 $a_i$（$1 \le a_i \le n$）。
定义 $\operatorname{diff}(u,v)$ 为从 $u$ 到 $v$ 的简单路径上不同颜色的个数。
对于一对节点 $(u,v)$，设 $p = \operatorname{lca}(u,v)$，定义

$$f(u,v) = \operatorname{diff}(u,p) \cdot \operatorname{diff}(v,p) $$

要求找出所有节点对 $(u,v)$（允许 $u=v$）中 $f(u,v)$ 的最大值。

思路分析

1. 固定 LCA 考虑
   对每个节点 $p$，考虑所有以 $p$ 为最近公共祖先的节点对 $(u,v)$。此时 $u$ 和 $v$ 分别位于 $p$ 的不同子树中（或其中一个等于 $p$）。
   若我们能快速得到对于固定的 $p$，其每个子节点 $c$ 的子树中 $\operatorname{diff}(p,x)$ 的最大值，那么取最大的两个值相乘即是以 $p$ 为 LCA 的最优答案。

2. 动态维护 $\operatorname{diff}(p,x)$
   我们采用 DFS 顺序处理：从根向下遍历，每当进入一个节点 $p$ 时，我们希望知道从 $p$ 到其子树内任意节点 $x$ 的路径上有多少种不同颜色。
   考虑在 DFS 过程中维护一个数组 $val[x]$，表示当前根节点 $p$ 到 $x$ 的不同颜色数。当我们从 $p$ 向下走到 $u$ 时，颜色 $a_u$ 会对 $u$ 的整个子树产生贡献（+1），但如果子树中还存在另一个同色节点 $v$，那么从 $p$ 到 $v$ 的路径上该颜色只能计一次，因此需要撤销 $u$ 的贡献。

3. 最近同色祖先
   定义 $\operatorname{up}[u]$ 为节点 $u$ 的祖先中颜色与 $u$ 相同且离 $u$ 最近的那个（若不存在则为 $0$）。
   那么当我们进入节点 $u$ 时，颜色 $a_u$ 的贡献应该只作用于 不包含任何 $\operatorname{up}[u]$ 的子节点 的子树区域。具体地：

   • 给 $u$ 的整个子树加 $1$；
   • 对于每个满足 $\operatorname{up}[v] = u$ 的节点 $v$，给 $v$ 的子树减 $1$。

   通过这种方式，线段树中维护的 $val[x]$ 正好等于 $\operatorname{diff}(p,x)$，其中 $p$ 是当前正在处理的节点（DFS 栈顶）。

4. 利用欧拉序转化为区间操作
   对树做一次 DFS 求出每个节点的进入时间 $\operatorname{in}[u]$ 和离开时间 $\operatorname{out}[u]$，则子树 $u$ 对应区间 $[\operatorname{in}[u], \operatorname{out}[u]]$。
   我们使用支持 区间加 和 区间最大值查询 的线段树。

算法步骤

步骤 1：计算 $\operatorname{up}[u]$ 和 $\operatorname{same}$ 列表

• 初始化全局数组 $\operatorname{last}[c]$ 表示当前 DFS 路径上颜色 $c$ 最后出现的节点。
• 在 DFS 进入节点 $u$ 时：
  • $\operatorname{up}[u] = \operatorname{last}[a_u]$；
  • 更新 $\operatorname{last}[a_u] = u$；
  • 递归处理子节点；
  • 回溯恢复 $\operatorname{last}[a_u]$。
• 若 $\operatorname{up}[u] \neq 0$，则将 $u$ 加入 $\operatorname{same}[\operatorname{up}[u]]$ 列表。

步骤 2：欧拉序预处理

• 第二次 DFS 得到每个节点的 $\operatorname{in}[u]$ 和 $\operatorname{out}[u]$。

步骤 3：第二次 DFS 计算答案

• 创建线段树，支持区间加和区间最大值。
• 执行 dfs2(u)：
  1. 区间 $[\operatorname{in}[u], \operatorname{out}[u]]$ 加 $1$（加入颜色 $a_u$ 的贡献）。
  2. 对 $\operatorname{same}[u]$ 中的每个节点 $v$，区间 $[\operatorname{in}[v], \operatorname{out}[v]]$ 减 $1$（撤销最近同色祖先的重复贡献）。
  3. 递归处理每个子节点 $v$，并在递归后查询区间 $[\operatorname{in}[v], \operatorname{out}[v]]$ 的最大值，存入列表 vals。
  4. 将节点 $u$ 自身对应的值 $1$ 也加入 vals（对应 $u=p$ 的情况）。
  5. 对 vals 降序排序，取前两个值 $m_1, m_2$，用 $m_1 \cdot m_2$ 更新全局答案 ans。
  6. 回溯：对 $\operatorname{same}[u]$ 中的每个 $v$，区间加 $1$；对 $u$ 的区间减 $1$。

步骤 4：输出答案

复杂度分析

• 每个节点在 dfs1 和 dfs2 中各被访问一次。
• 线段树每次区间修改和查询均为 $O(\log n)$。
• 总时间复杂度 $O(n \log n)$，空间复杂度 $O(n)$。
  由于所有测试用例的 $n$ 之和不超过 $3\times 10^5$，完全可行。

代码实现

#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
using namespace std;

const int MAXN = 300005;

int n;
vector<int> children[MAXN];
int a[MAXN];
int up[MAXN]; // nearest ancestor with same color
int last[MAXN]; // last occurrence of each color in current path
vector<int> same[MAXN]; // vertices for which this vertex is the nearest same-color ancestor

int in[MAXN], out[MAXN], timer;
int euler[MAXN]; // not really needed

// Segment tree for range add and range max
struct SegTree {
    int n;
    vector<int> maxv, lazy;
    SegTree(int _n) : n(_n), maxv(4 * _n), lazy(4 * _n) {}
    void push(int idx) {
        if (lazy[idx]) {
            maxv[idx*2] += lazy[idx];
            maxv[idx*2+1] += lazy[idx];
            lazy[idx*2] += lazy[idx];
            lazy[idx*2+1] += lazy[idx];
            lazy[idx] = 0;
        }
    }
    void update(int l, int r, int val, int idx, int L, int R) {
        if (l <= L && R <= r) {
            maxv[idx] += val;
            lazy[idx] += val;
            return;
        }
        push(idx);
        int mid = (L + R) / 2;
        if (l <= mid) update(l, r, val, idx*2, L, mid);
        if (r > mid) update(l, r, val, idx*2+1, mid+1, R);
        maxv[idx] = max(maxv[idx*2], maxv[idx*2+1]);
    }
    int query(int l, int r, int idx, int L, int R) {
        if (l <= L && R <= r) return maxv[idx];
        push(idx);
        int mid = (L + R) / 2;
        int res = 0;
        if (l <= mid) res = max(res, query(l, r, idx*2, L, mid));
        if (r > mid) res = max(res, query(l, r, idx*2+1, mid+1, R));
        return res;
    }
    void update(int l, int r, int val) { update(l, r, val, 1, 1, n); }
    int query(int l, int r) { return query(l, r, 1, 1, n); }
};

void dfs1(int u) {
    up[u] = last[a[u]];
    last[a[u]] = u;
    for (int v : children[u]) {
        dfs1(v);
    }
    last[a[u]] = up[u];
    if (up[u] != 0) {
        same[up[u]].push_back(u);
    }
}

void dfs_euler(int u) {
    in[u] = ++timer;
    for (int v : children[u]) {
        dfs_euler(v);
    }
    out[u] = timer;
}

long long ans;
SegTree *seg;

void dfs2(int u) {
    // add current vertex's color
    seg->update(in[u], out[u], 1);
    // subtract for vertices that have u as nearest same-color ancestor
    for (int v : same[u]) {
        seg->update(in[v], out[v], -1);
    }
    vector<int> vals;
    // process children
    for (int v : children[u]) {
        dfs2(v);
        int cur = seg->query(in[v], out[v]);
        vals.push_back(cur);
    }
    // add u itself (value = 1)
    vals.push_back(1);
    // find top two
    sort(vals.rbegin(), vals.rend());
    if (vals.size() >= 2) {
        ans = max(ans, 1LL * vals[0] * vals[1]);
    } else if (vals.size() == 1) {
        ans = max(ans, 1LL * vals[0]);
    }
    // rollback
    for (int v : same[u]) {
        seg->update(in[v], out[v], 1);
    }
    seg->update(in[u], out[u], -1);
}

void solve() {
    cin >> n;
    for (int i = 1; i <= n; i++) {
        children[i].clear();
        same[i].clear();
    }
    for (int i = 2; i <= n; i++) {
        int p; cin >> p;
        children[p].push_back(i);
    }
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    // initialize last array with 0
    memset(last, 0, sizeof(last));
    dfs1(1);
    timer = 0;
    dfs_euler(1);
    seg = new SegTree(n);
    ans = 1;
    dfs2(1);
    cout << ans << "\n";
    delete seg;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}

总结

本题的核心思想是 在 DFS 过程中动态维护从当前节点到子树内任意节点的不同颜色数，并利用 最近同色祖先 来避免重复计数。通过欧拉序将子树转化为区间，使用线段树支持区间加和区间最大值，最终在 $O(n \log n)$ 时间内求出答案。该解法巧妙地将树上的路径计数问题转化为区间操作，值得深入理解。