【BZOJ 2783】【JLOI 2012】树【STL-set】

题意

第一行是两个整数N和S，其中N是树的节点数。
第二行是N个正整数，第i个整数表示节点i的正整数。
接下来的N-1行每行是2个整数x和y，表示y是x的儿子。

输出格式：

输出路径节点总和为S的路径数量。

题解

　　因为求的路径都是从在包含根的一条链上，而且是从上往下。所以可以考虑直接dfs。记录元素前缀和。
　　比如搜到x，在set里面查询有没有s[x]-k，有的话ans++，然后继续往下dfs时插入s[x]。
　　最开始要插入0，同时在离开点x时要记住删除那个s[x]。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<set>

#define N 100010
using namespace std;

multiset<int> q;
struct node{int to,next;}e[N*2];
int head[N],m;
int ans,a[N],n,s,sum[N];

void add_edge(int from,int to) {
    e[++m].next = head[from];
    head[from] = m;
    e[m].to = to;
}

void dfs(int v) {
    if(q.find(sum[v]-s) != q.end()) ans++;
    q.insert(sum[v]);
    for(int i = head[v];i;i=e[i].next) {
        sum[e[i].to] = sum[v] + a[e[i].to];
        dfs(e[i].to);
    }
    q.erase(q.find(sum[v]));
}

int main()
{
    int u,v;
    scanf("%d%d",&n,&s);
    for(int i = 1;i <= n;i++) scanf("%d",&a[i]);
    m = 0;
    for(int i = 1;i < n;i++) {
        scanf("%d%d",&u,&v);
        add_edge(u,v);
    }

    q.insert(0); ans = 0;
    sum[1] = a[1]; dfs(1);
    printf("%d\n",ans);
    return 0;
}