HTR003E NOI树 题解

题目 大意：

给出一颗 $n$ 个结点的树，每个结点上有 N，O，I 三个字母其中一个。一个 NOI 三元组”定义为：三个分别为 N，O，I 的结点，且 O 结点在 N 结点和 I 结点的路径上。求 NOI 三元组的个数，对 $10^9+7$ 取模。

考虑钦定一个点为根，记录以每个点为根的子树内 N 和 I 的个数，这个通过 dfs 就可以求出。在记录总共有多少个 N 和 I， 这样就可以求出这个点子树外的 N 和 I 的个数。接下来枚举中间点 O ，扫一遍即可。

具体的，每次枚举到一个 O 时，遍历所有子树，记录之前遍历到的所有子树 N和 O 的数量，那么这个与前面的子树内的点匹配的贡献就珂以直接相乘算出。与后面子树内的点的贡献将在后面计算，这样就可以做到不重不漏。

不懂的话可以看代码理解 Qwq。

CODE :

#include <bits/stdc++.h>
using namespace std;
#define MxN 100010
#define p 1000000007

int n, a[MxN], hd[MxN], res = 0, tot = 0, N[MxN], I[MxN], sN = 0, sI = 0, f[MxN];
struct node {
    int y, to;
    node(int y = 0, int to = 0) : y(y), to(to) {}
}e[MxN * 2];

void add(int x, int y) {
    e[++ tot] = node(y, hd[x]), hd[x] = tot;
}

void dfs(int x, int fa) {
    f[x] = fa;
    if (a[x] == 1) N[x] = 1;
    else if (a[x]) I[x] = 1;
    for (int i = hd[x]; i; i = e[i].to) {
        int y = e[i].y;
        if (y == fa) continue;
        dfs(y, x);
        N[x] += N[y], I[x] += I[y];
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin >> n;
    for (int i = 1; i <= n; i ++) {
        char c; cin >> c;
        if (c == 'O') a[i] = 0;
        if (c == 'N') a[i] = 1, sN ++;
        if (c == 'I') a[i] = 2, sI ++;
    } 
    for (int i = 1, x, y; i < n; i ++) {
        cin >> x >> y, add(x, y), add(y, x);
    }
    dfs(1, 0);
    for (int i = 1; i <= n; i ++) {
        if (a[i]) continue;
        int tN = 0, tI = 0;
        for (int j = hd[i]; j; j = e[j].to) {
            int y = e[j].y;
            if (y == f[i]) continue;
            (res += N[y] * tI) %= p, tI += I[y];
            (res += I[y] * tN) %= p, tN += N[y];
        }
        (res += (sN - tN) * tI) %= p;
        (res += (sI - tI) * tN) %= p;
    }
    cout << res << endl;
    return 0;
}