2019牛客国庆集训派对day3 买一送一

题目链接：

题意：有n个点，n-1条单向边，每个点都销售一类商品

问从点1开始走，买第一样商品类型为x,买第二样商品类型为y，问不同有序对<x,y>的数量

解法：

col[i]表示这个点的商品类型

last[col[i]]表示从1到点i过程中，点i的商品类型上次出现的时候的点的父亲

vis[col[i]]表示从1到点i过程中，点i的商品类型经过次数

num[i]表示从1到点i过程中不同商品类型数量和

每次扫到新的点v时，(u为v的父亲) num[u] - num[last[col[v]]]就是需要更新的对数

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int M = 1e5 + 10;
const int mod = 1e9 + 7;
int num[M], vis[M], anss[M], ans, cnt, col[M], head[M], last[M];
struct node{
    int next, to;
}edge[M];
void add_edge(int u, int v) {
    edge[++cnt].next = head[u];
    edge[cnt].to = v;
    head[u] = cnt;
}
void dfs(int u, int sum, int ans) {
    for(int i = head[u]; i; i = edge[i].next) {
        int v = edge[i].to;
        int t = 0;
        int lastt = last[col[v]];
        vis[col[v]]++;
        if(vis[col[v]] == 1) num[v] = num[u] + 1;
        else num[v] = num[u];
        t = num[u] - num[last[col[v]]];
        last[col[v]] = u;
        //printf("%d %d %d\n", v, col[v], last[col[v]]);
        anss[v] = ans + t;
        if(vis[col[v]] == 1) dfs(v, sum + 1, ans + t);
        else dfs(v, sum + 1, ans + t);
        vis[col[v]]--;
        last[col[v]] = lastt;
    }
}
int main(){
    int n;
    while(~scanf("%d", &n)){
        cnt = 0;
        memset(head, 0, sizeof(head));
        for(int i = 1; i <= n; i++) num[i] = 0, last[i] = 0, vis[i] = 0, anss[i] = 0;
        for(int i = 2; i <= n; i++) {
            int u;
            scanf("%d", &u);
            add_edge(u, i);
        }
        for(int i = 1; i <= n; i++)
            scanf("%d", &col[i]);
        vis[col[1]]++;
        num[1] = 1;
        dfs(1, 1, 0);
//        for(int i = 1; i <= n; i++) {
//            printf("%d ", num[i]);
//        }
//        printf("\n");
        for(int i = 2; i <= n; i++) {
            printf("%d\n", anss[i]);
        }
    }
    return 0;
}
/*
3
1 2
1 2 3
3
1 1
1 2 3
4
1 2 3
1 3 2 3
7
1 1 3 2 4 2
3 3 3 4 5 3 3
7
1 1 3 2 4 2
2 3 3 4 5 3 3
*/

View Code