[luoguP3252] [JLOI2012]树（DP）

传送门

 

树上前缀和。

在树上找一条权值和为 s 的链，其中这个链上的点按深度递增（递减）（不同）

dfs

每搜到一个点求它的前缀和 sum[x]，放入 set 中。

在 set 中找 sum[x] - s 的点，如果有，ans++

退出 dfs 的时候再把 sum[x] 从 set 中删除

因为每个点权都是正整数，所以 set 中没有重复元素。

同时也是单调递增，所以简单些不用 set，开个数组再 lower_bound 也行。

 

——代码

#include <set>
#include <cstdio>
#include <cstring>
#include <iostream>

const int MAXN = 100001;
int n, s, cnt, ans;
int a[MAXN], head[MAXN], to[MAXN << 1], next[MAXN << 1], f[MAXN], sum[MAXN];
std::set <int> S;

inline int read()
{
    int x = 0, f = 1;
    char ch = getchar();
    for(; !isdigit(ch); ch = getchar()) if(ch == '-') f = -1;
    for(; isdigit(ch); ch = getchar()) x = (x << 1) + (x << 3) + ch - '0';
    return x * f;
}

inline void add(int x, int y)
{
    to[cnt] = y;
    next[cnt] = head[x];
    head[x] = cnt++;
}

inline void dfs(int u)
{
    sum[u] = sum[f[u]] + a[u];
    S.insert(sum[u]);
    if(S.count(sum[u] - s)) ans++;
    for(int i = head[u]; i ^ -1; i = next[i]) dfs(to[i]);    
    S.erase(sum[u]);
}

int main()
{
    int i, x, y;
    n = read();
    s = read();
    memset(head, -1, sizeof(head));
    for(i = 1; i <= n; i++) a[i] = read();
    for(i = 1; i < n; i++)
    {
        x = read();
        y = read();
        f[y] = x;
        add(x, y);
    }
    S.insert(0);
    dfs(1);
    printf("%d\n", ans);
    return 0;
}