#5456. 「ROI 2012 Day 1」病毒与杀毒软件 题解

问题分析

我们有两个树结构：

• 正式树：公司公开的管理层级
• 秘密树：员工私下组成的病毒开发层级

稳定对定义为：员工 $A$ 在两棵树中都是 $B$ 的祖先。

我们需要统计这样的 $(A,B)$ 对的数量，其中 $A \neq B$。

关键观察

1. 树结构性质：在两棵树中，如果 $A$ 是 $B$ 的祖先，则 $A$ 在 DFS 序中出现在 $B$ 之前，且 $B$ 在 $A$ 的子树中。

2. 问题转化：对于每个员工 $B$，统计在两棵树中都是 $B$ 的祖先的员工 $A$ 的数量，然后对所有 $B$ 求和。

算法思路

方法：DFS 序 + 树状数组/线段树

步骤：

1. 预处理两棵树的 DFS 序

   • 对每棵树进行 DFS，记录每个节点的进入时间 $in[u]$ 和离开时间 $out[u]$
   • 在 DFS 序中，$u$ 的子树对应区间 $[in[u], out[u]]$

2. 利用正式树 DFS 序排序

   • 按照正式树的 DFS 进入时间排序所有节点

3. 在秘密树上使用数据结构

   • 按正式树 DFS 序顺序处理节点
   • 对于当前节点 $u$，在秘密树中查询 $u$ 的祖先数量
   • 然后将 $u$ 在秘密树中的位置加入数据结构

具体实现

设：

• $in_1[u], out_1[u]$：正式树的 DFS 序
• $in_2[u], out_2[u]$：秘密树的 DFS 序

算法流程：

1. 对正式树进行 DFS，得到 $in_1[u], out_1[u]$
2. 对秘密树进行 DFS，得到 $in_2[u], out_2[u]$
3. 创建数组 $order$，按 $in_1[u]$ 排序所有节点
4. 初始化树状数组（大小为秘密树的 DFS 序范围）
5. 按 $order$ 顺序处理每个节点 $u$：
   • 查询树状数组中位置 $in_2[u]$ 的前缀和 → 这是在秘密树中 $u$ 的祖先数量
   • 将答案加上这个数量
   • 在树状数组的位置 $in_2[u]$ 处加 1
   • 在树状数组的位置 $out_2[u]+1$ 处减 1（DFS 序区间更新）

复杂度分析

• 时间复杂度：$O(N \log N)$
  • 两次 DFS：$O(N)$
  • 排序：$O(N \log N)$
  • 树状数组操作：$O(N \log N)$
• 空间复杂度：$O(N)$

代码实现

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

const int MAXN = 100005;

vector<int> tree1[MAXN], tree2[MAXN];
int in1[MAXN], out1[MAXN], in2[MAXN], out2[MAXN];
int timer1, timer2;

// 树状数组
int bit[MAXN * 2];

void update(int idx, int val) {
    for (; idx < MAXN * 2; idx += idx & -idx)
        bit[idx] += val;
}

int query(int idx) {
    int sum = 0;
    for (; idx > 0; idx -= idx & -idx)
        sum += bit[idx];
    return sum;
}

// DFS 遍历树
void dfs1(int u) {
    in1[u] = ++timer1;
    for (int v : tree1[u]) {
        dfs1(v);
    }
    out1[u] = timer1;
}

void dfs2(int u) {
    in2[u] = ++timer2;
    for (int v : tree2[u]) {
        dfs2(v);
    }
    out2[u] = timer2;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    int N;
    cin >> N;
    
    // 读取正式树
    int root1 = -1;
    for (int i = 1; i <= N; i++) {
        int a;
        cin >> a;
        if (a == 0) {
            root1 = i;
        } else {
            tree1[a].push_back(i);
        }
    }
    
    // 读取秘密树  
    int root2 = -1;
    for (int i = 1; i <= N; i++) {
        int b;
        cin >> b;
        if (b == 0) {
            root2 = i;
        } else {
            tree2[b].push_back(i);
        }
    }
    
    // DFS 遍历
    dfs1(root1);
    dfs2(root2);
    
    // 按正式树 DFS 进入时间排序
    vector<pair<int, int>> nodes; // (in1[u], u)
    for (int i = 1; i <= N; i++) {
        nodes.push_back({in1[i], i});
    }
    sort(nodes.begin(), nodes.end());
    
    long long ans = 0;
    
    // 处理节点
    for (auto [_, u] : nodes) {
        // 查询在秘密树中 u 的祖先数量
        ans += query(in2[u]);
        
        // 更新树状数组（标记 u 在秘密树中的子树范围）
        update(in2[u], 1);
        update(out2[u] + 1, -1);
    }
    
    cout << ans << endl;
    
    return 0;
}

正确性证明

为什么这样工作：

1. 按正式树 DFS 序处理，保证处理 $u$ 时，所有在正式树中是 $u$ 祖先的节点都已处理
2. 树状数组维护了在秘密树中当前已处理节点的子树信息
3. 查询 $in_2[u]$ 得到在秘密树中 $u$ 的祖先数量

满足稳定对条件：

• 如果 $A$ 在正式树中是 $B$ 的祖先，则 $in_1[A] < in_1[B]$，所以 $A$ 先于 $B$ 处理
• 如果 $A$ 在秘密树中是 $B$ 的祖先，则 $in_2[A] < in_2[B] \leq out_2[A]$，查询时能正确计数