Fixed-Length Paths I

#ifdef LOCAL
#include <iostream>
#include <fstream>
#include <iomanip>
#include <vector>
#include <queue>
#include <stack>
#include <set>
#include <map>
#else
#include <bits/stdc++.h>
#define cerr if (false) cerr
#define endl '\n'
#endif

#define fi first
#define se second

#define sz(a) ((int)(a).size())
#define all(a) (a).begin(), (a).end()

#define lsb(x) (x & (-x))

using namespace std;

template <typename T>
bool ckmax(T &a, T b) { return a < b ? a = b, true : false; }
template <typename T>
bool ckmin(T &a, T b) { return a > b ? a = b, true : false; }

using ll = long long;
using pii = pair<int, int>;

const int NMAX = 2e5+5;
//const int INF = 1e9+5;

int n, k;
vector<int> g[NMAX];
int blocked[NMAX];
int sz[NMAX];
int fr[NMAX];
int maxd;
ll ans;

void read() {
   cin >> n >> k;
   for (int i = 1, x, y; i < n; i++) {
      cin >> x >> y;
      g[x].push_back(y);
      g[y].push_back(x);
   }
}

void dfs_size(int x, int p = 0) {
   sz[x] = 1;
   for (auto &u : g[x]) {
      if (u == p || blocked[u]) continue;
      dfs_size(u, x);
      sz[x] += sz[u];
   }
}

int find_centroid(int x, int cnt, int p = 0) {
   for (auto &u : g[x]) {
      if (u == p || blocked[u]) continue;
      if (sz[u] << 1 >= cnt)
         return find_centroid(u, cnt, x);
   }
   return x;
}

void dfs_count(int x, int p, int d) {
   maxd = max(maxd, d);
   ans += fr[k - d];
   if (k <= d) return;
   for (auto &u : g[x]) {
      if (u == p || blocked[u]) continue;
      dfs_count(u, x, d + 1);
   }
}

void dfs_mark(int x, int p, int d) {
   fr[d]++;
   for (auto &u : g[x]) {
      if (u == p || blocked[u]) continue;
      dfs_mark(u, x, d + 1);
   }
}

void centroid_decomp(int x = 1) {
   dfs_size(x);

   x = find_centroid(x, sz[x]);
   blocked[x] = true;

   fr[0] = 1;
   maxd = 0;

   for (auto &u : g[x]) {
      if (blocked[u]) continue;
      dfs_count(u, x, 1);
      dfs_mark(u, x, 1);
   }

   memset(fr, 0, (maxd + 1) * sizeof(int));

   for (auto &u : g[x]) {
      if (blocked[u]) continue;
      centroid_decomp(u);
   }
}

ll solve() {
   centroid_decomp();
   return ans;
}

signed main() {
#ifdef LOCAL
   freopen("input.txt", "r", stdin);
#endif
   ios_base::sync_with_stdio(false);
   cin.tie(0);
   cout.tie(0);

   int t = 1;
//   cin >> t;

   while (t--) {
      read();
      cout << solve() << endl;
   }

   return 0;
}