Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include<bits/stdc++.h>
- #define ll long long
- #define pb push_back
- #define f first
- #define s second
- #define mp make_pair
- using namespace std;
- const int MAXN = 1e5 + 100;
- const int MOD = 1e9 + 7;
- int dp[MAXN], dpH[MAXN];
- vector <int> g[MAXN];
- bool removed[MAXN];
- ll ans;
- map <int, int> allDist;
- int n, k;
- int sizeComp(int x, int pr = -1){
- dp[x] = 1;
- for (int i = 0; i < g[x].size(); i++){
- int to = g[x][i];
- if (to == pr)
- continue;
- if (!removed[to]){
- dp[x] += sizeComp(to, x);
- }
- }
- return dp[x];
- }
- void findDist(int x, int h = 1, int pr = -1){
- dpH[x] = h;
- ans += allDist[k - dpH[x]];
- for (int i = 0; i < g[x].size(); i++){
- int to = g[x][i];
- if (to == pr)
- continue;
- if (!removed[to]){
- findDist(to, h + 1, x);
- }
- }
- }
- void addDist(int x, int pr = -1){
- allDist[dpH[x]]++;
- //cout << "x = " << x << ' ' << "dpH[x] = " << dpH[x] << endl;
- for (int i = 0; i < g[x].size(); i++){
- int to = g[x][i];
- if (to == pr)
- continue;
- if (!removed[to]){
- addDist(to, x);
- }
- }
- }
- int findC(int x, int sz, int pr = -1){
- for (int i = 0; i < g[x].size(); i++){
- int to = g[x][i];
- if (to == pr)
- continue;
- if (!removed[to] && dp[to] * 2 > sz){
- return findC(to, sz, x);
- }
- }
- return x;
- }
- void calc(int c){
- //cout << "c = " << c << endl;
- for (int i = 0; i < g[c].size(); i++){
- int to = g[c][i];
- if (!removed[to]){
- findDist(to);
- addDist(to);
- }
- }
- ans += allDist[k];
- }
- void solve(int x, int cnt = 0){
- int sz = sizeComp(x);
- int c = findC(x, sz);
- removed[c] = 1;
- //ans[c] = cnt;
- allDist.clear();
- calc(c);
- for (int i = 0; i < g[c].size(); i++){
- int to = g[c][i];
- if (!removed[to]){
- solve(to, cnt + 1);
- }
- }
- }
- int main(){
- ios::sync_with_stdio(0);
- cin.tie(0);
- cout.tie(0);
- cin >> n >> k;
- for (int i = 1; i < n; i++){
- int x, y;
- cin >> x >> y;
- g[x].pb(y);
- g[y].pb(x);
- }
- solve(1);
- cout << ans;
- }
Add Comment
Please, Sign In to add comment