[牛客] [#1108 J] [树形结构] 买一送一

2019牛客国庆集训派对day3

链接：https://ac.nowcoder.com/acm/contest/1108/J
来源：牛客网

题意

ICPCCamp 有 $n$ 个商店，用 $1,2,...,n$ 编号。对于任意 $i > 1$，有从商店 $p_i$​ 到 $i$ 的单向道路。
同时，商店 $i$ 出售类型为 $a_i$​ 的商品。
Bobo 从商店 $1$ 出发前往商店 $i$。他要在两个不同的商店购买商品（包括商店 $1$ 和$i$）。设他先购买的商品类型是 $x$，后购买的商品类型是 $y$，他用 $f_i$​ 表示不同的有序对 $\langle x, y \rangle$ 的数量。
求出 $f_2, f_3, \dots, f_n$​ 的值。

思路

由于每个点只有一个入度，所以可以转化为根为$1$的树形结构。

那么问题就转化为，从根$1$ 到 $i$ 的路径上不同有序对的个数。

 

// #pragma GCC optimize(2)
// #pragma GCC optimize(3)
// #pragma GCC optimize(4)
#include <algorithm>
#include <iterator>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <iomanip>
#include <bitset>
#include <cctype>
#include <cstdio>
#include <string>
#include <vector>
#include <stack>
#include <cmath>
#include <queue>
#include <list>
#include <map>
#include <set>
#include <cassert>
//#include <unordered_set>
//#include <unordered_map>
// #include<bits/extc++.h>
// using namespace __gnu_pbds;
using namespace std;
#define pb push_back
#define fi first
#define se second
#define debug(x) cerr << #x << " := " << x << endl;
#define bug cerr << "-----------------------" << endl;
#define FOR(a, b, c) for (int a = b; a <= c; ++a)

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int inf = 0x3f3f3f3f;
const ll inff = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 + 7;

template <typename T>
inline T read(T &x) {
    x = 0;
    int f = 0;
    char ch = getchar();
    while (ch < '0' || ch > '9') f |= (ch == '-'), ch = getchar();
    while (ch >= '0' && ch <= '9') x = x * 10 + ch - '0', ch = getchar();
    return x = f ? -x : x;
}

/**********showtime************/
const int maxn = 1e5 + 9;
int p[maxn], a[maxn];
ll f[maxn];
int sum, kind[maxn], cnt[maxn];
vector<int> mp[maxn];

ll add(int col) {
    ll res = 0;
    res = sum - kind[col];
    if (cnt[col] == 1)
        res++;
    else if (cnt[col] == 0)
        sum++;
    cnt[col]++;
    kind[col] = sum;
    return res;
}
void del(int col) {
    cnt[col]--;
    if (cnt[col] == 0) sum--;
}

void dfs(int u, int fa, int dp) {
    f[u] = f[fa];
    int yk = kind[a[u]];
    f[u] += add(a[u]);

    for (int v : mp[u]) {
        dfs(v, u, dp + 1);
    }

    del(a[u]);

    kind[a[u]] = yk;
}
int main() {
    int n;
    while (~scanf("%d", &n)) {
        for (int i = 1; i <= n; i++) mp[i].clear(), f[i] = 0, kind[i] = 0;
        sum = 0;

        for (int i = 2; i <= n; i++) {
            scanf("%d", &p[i]);
            mp[p[i]].pb(i);
        }
        for (int i = 1; i <= n; i++) scanf("%d", &a[i]);

        dfs(1, 1, 1);

        for (int i = 2; i <= n; i++) printf("%lld\n", f[i]);
    }
    return 0;
}