fork download
copy 
  1. #include <iostream>     // Input/output stream objects
  2. #include <fstream>      // File stream objects
  3. #include <sstream>      // String stream objects
  4. #include <iomanip>      // Input/output manipulators
  5. #include <string>       // String class and functions
  6. #include <vector>       // Dynamic array
  7. #include <list>         // Doubly linked list
  8. #include <set>          // Set container
  9. #include <map>          // Map container
  10. #include <queue>        // Queue container
  11. #include <stack>        // Stack container
  12. #include <algorithm>    // Algorithms on sequences (e.g., sort, find)
  13. #include <cmath>        // Mathematical functions
  14. #include <ctime>        // Date and time functions
  15. #include <cstdlib>      // General purpose functions (e.g., memory management)
  16. #include <cstring>      // C-style string functions
  17. #include <cctype>       // Character classification functions
  18. #include <cassert>      // Assert function for debugging
  19. #include <exception>    // Standard exceptions
  20. #include <functional>   // Function objects
  21. #include <iterator>     // Iterator classes
  22. #include <limits>       // Numeric limits
  23. #include <locale>       // Localization and internationalization
  24. #include <numeric>      // Numeric operations (e.g., accumulate)
  25. #include <random>       // Random number generators
  26. #include <stdexcept>    // Standard exception classes
  27. #include <typeinfo>     // Runtime type information
  28. #include <utility>      // Utility components (e.g., std::pair)
  29. #include <complex>      // Complex Module
  30. #define int long long 
  31.  
  32. using namespace std;
  33. using ll = long long;
  34. const int N = 2e5 + 5;
  35.  
  36. int n, k, ret, blocked[N], sz[N];
  37. vector<int> adj[N];
  38.  
  39. void pre(int u, int p) {
  40.     sz[u] = 1;
  41.     for (auto& v : adj[u]) {
  42.         if (v == p || blocked[v]) continue;
  43.         pre(v, u);
  44.         sz[u] += sz[v];
  45.     }
  46. }
  47.  
  48. int find_centroid(int u, int p, int tot) {
  49.     for (auto& v : adj[u]) {
  50.         if (v == p || blocked[v] || 2 * sz[v] <= tot) continue;
  51.         return find_centroid(v, u, tot);
  52.     }
  53.     return u;
  54. }
  55.  
  56. int calc(int u, int p, int dx) {
  57.     vector<int> d(n + 1, n + 5);
  58.     d[u] = dx;
  59.     queue<int> q;
  60.     q.push(u);
  61.     while(!q.empty()) {
  62.         int a = q.front();
  63.         q.pop();
  64.         for(auto& b : adj[a]) {
  65.             if (b != p && !blocked[b] && d[b] == n + 5) {
  66.                 d[b] = d[a] + 1;
  67.                 q.push(b);
  68.             }
  69.         }
  70.     }
  71.     int ans = 0;
  72.     vector<int> freq(n + 1, 0);
  73.     for (int i = 1; i <= n; i++) {
  74.         if (d[i] != n + 5) freq[d[i]]++;
  75.     }
  76.     for (int i = 0, j = k; i <= j; i++, j--) {
  77.         if (i == j) 
  78.             ans += freq[i] * (freq[i] - 1)/2;
  79.         else
  80.             ans += freq[i] * freq[j];
  81.     }
  82.     return ans;
  83. }
  84.  
  85. void solve(int u) {
  86.     pre(u, 0);
  87.  
  88.     int cent = find_centroid(u, 0, sz[u]);
  89.  
  90.     ret += calc(cent, 0, 0);
  91.     for (auto& v : adj[cent]) {
  92.         if (!blocked[v]) ret -= calc(v, cent, 1);
  93.     }
  94.  
  95.     blocked[cent] = 1;
  96.  
  97.     for (auto& v : adj[cent]) {
  98.         if (!blocked[v]) {
  99.             solve(v);
  100.         }
  101.     }
  102. }
  103.  
  104. int32_t main() {
  105.     ios::sync_with_stdio(false);
  106.     cin.tie(nullptr), cout.tie(nullptr);
  107.  
  108.     cin >> n >> k;
  109.     for (int i = 1; i < n;i ++) {
  110.         int u, v;
  111.         cin >> u >> v;
  112.         adj[u].push_back(v);
  113.         adj[v].push_back(u);
  114.     }
  115.  
  116.     solve(1);
  117.  
  118.     cout << ret << endl;
  119.  
  120.     return 0;
  121. }
  122.  
Success #stdin #stdout 0.01s 10548KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
0
https://ideone.com/07QUzW
language:
C++14 (gcc 8.3)
created:
3 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!