Skip to content

11591: 【原1591】"Count On Tree"

题目

题目描述

author: Crystal 原OJ链接:https://acm.sjtu.edu.cn/OnlineJudge-old/problem/1591

Description

Crystal家有一棵树。树上有n个节点,编号由1到n(1号点是这棵树的根),两点之间距离为1当且仅当它们直接相连。每个点都有各自的权值,第i号节点的权值为value_i。Crystal现在指着编号为x的点问,在以点x为根的子树中,与点x距离大于等于k的所有点的点权和是多少。

Input Format

第1行两个整数n,Q,分别表示树上点的个数和Crystal有Q个问题。

第2行,n个整数,分别表示1至n号点的点权。

接下来的n - 1行,每行两个整数u,v,表示编号为u的点与编号为v的点直接相连。

接下来Q行,每行两个整数x, k,表示询问在以点x为根的子树中,与点x距离大于等于为k的所有点的点权和是多少。

Output Format

Q行,每行一个整数,表示对第i个询问的回答。

Sample Input

5 3
1 1 1 1 1
1 2
1 3
3 4
4 5
1 3
1 2
1 1

Sample Output

1
2
4

Limits

  • 对于30%的数据,保证n <= 1000, k < 1, Q <= 1000。
  • 对于60%的数据,保证n <= 1000, k < 1000, Q <= 1000
  • 对于80%的数据,保证n <= 1000, k < 1000, Q <= 1000000;
  • 对于最后20%的数据,保证n <= 50000, k < 100, Q <= 1000000;
  • 对于100%的数据,保证所有输入数据均为非负整数,且在int范围内。

FineArtz's solution

/* Count On Tree*/
#include <iostream>
#include <vector>
#include <queue>
using namespace std;

class Node{
public:
    int father = 0;
    vector<int> edge;
    long long w = 0, sumw = 0;
    int depth = 0;
    int inT = 0, outT = 0;
};
Node a[50005];
vector<int> dep[50105];

void makeTree(){
    queue<int> q;
    q.push(1);
    int now, next;
    while (!q.empty()){
        now = q.front();
        q.pop();
        for (int i : a[now].edge){
            if (i != a[now].father){
                a[i].father = now;
                if (a[i].edge.size() != 1)
                    q.push(i);
            }
        }
    }
}
/*void print(int root){
    if (root == 0) return;
    cout << root << ' ' << a[root].sumw << endl;
    print(a[root].lchild);
    print(a[root].rchild);
}*/
long long calcw(int root){
    a[root].sumw = a[root].w;
    for (auto i : a[root].edge){
        if (i != a[root].father){
            if (a[i].edge.size() != 1)
                a[root].sumw += calcw(i);
            else{
                a[i].sumw = a[i].w;
                a[root].sumw += a[i].w;
            }
        }
    }
    return a[root].sumw;
}
void calcDepth(int root, int d){
    if (root == 0) return;
    a[root].depth = d;
    dep[d].push_back(root);
    for (auto i : a[root].edge){
        if (i != a[root].father){
            calcDepth(i, d + 1);
        }
    }
}

int timeStamp = 1;
void makeStamp(int root){
    a[root].inT = timeStamp;
    for (auto i : a[root].edge){
        if (i != a[root].father){
            ++timeStamp;
            makeStamp(i);
        }
    }
    a[root].outT = timeStamp;
}

long long query(int root, int h){
    long long ret = 0;
    for (auto i : dep[a[root].depth + h]){
        if (a[i].inT >= a[root].inT && a[i].outT <= a[root].outT)
            ret += a[i].sumw;
    }
    return ret;
}
int main(){
    ios::sync_with_stdio(false);
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    int n, q, u, v;
    cin >> n >> q;
    for (int i = 1; i <= n; ++i)
        cin >> a[i].w;
    for (int i = 1; i < n; ++i){
        cin >> u >> v;
        a[u].edge.push_back(v);
        a[v].edge.push_back(u);
    }
    makeTree();
    calcw(1);
    //print(1);
    calcDepth(1, 1);
    makeStamp(1);
    int x, k;
    for (int i = 1; i <= q; ++i){
        cin >> x >> k;
        cout << query(x, k) << '\n';
    }
    return 0;
}