Submission #587714

#	Time	Username	Problem	Language	Result	Execution time	Memory
587714		MilosMilutinovic	Janjetina (COCI21_janjetina)	C++14	110 / 110	279 ms	19964 KiB

Main

/**
 *    author:  wxhtzdy
 *    created: 02.07.2022 11:01:49
**/
#include <bits/stdc++.h>

using namespace std;
             
template <typename T>
class fenwick {
  public:
  vector<T> fenw;
  int n;
  fenwick(int _n) : n(_n) {
    fenw.resize(n);
  }
  void modify(int x, T v) {
    while (x < n) {
      fenw[x] += v;
      x |= (x + 1);
    }
  }
  T get(int x) {
    T v{};
    while (x >= 0) {
      v += fenw[x];
      x = (x & (x + 1)) - 1;
    }
    return v;
  }
};

long long Calc(vector<pair<int, int>> v, int k) {
  int n = (int) v.size();
  sort(v.begin(), v.end());
  fenwick<int> fenw(n);
  long long ans = 0;
  for (int i = 0; i < n; i++) {
    int d = min(v[i].first - v[i].second - k, n - 1);
    ans += fenw.get(d);
    fenw.modify(v[i].second, +1);
  }
  return ans;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);  
  int n, k;
  cin >> n >> k;
  vector<vector<pair<int, int>>> g(n);
  for (int i = 0; i < n - 1; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    --u; --v;
    g[u].emplace_back(v, w);
    g[v].emplace_back(u, w);
  }
  long long ans = 0;
  vector<int> sz(n);
  vector<bool> was(n);
  function<void(int, int)> Dfs = [&](int v, int pr) {
    sz[v] = 1;
    for (auto& e : g[v]) {
      int u = e.first;
      if (was[u] || u == pr) {
        continue;
      }
      Dfs(u, v);
      sz[v] += sz[u];
    }
  };
  function<int(int, int, int)> Find = [&](int v, int pr, int n) {
    for (auto& e : g[v]) {
      int u = e.first;
      if (was[u] || u == pr || sz[u] * 2 < n) {
        continue;
      }
      return Find(u, v, n);
    }
    return v;
  };                                          
  vector<pair<int, int>> vec;
  function<void(int, int, int, int)> Go = [&](int v, int pr, int mx, int len) {
    vec.emplace_back(mx, len);
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (!was[u] && u != pr) {
        Go(u, v, max(mx, w), len + 1);
      }
    }
  };
  function<void(int)> Solve = [&](int v) {
    Dfs(v, v);
    v = Find(v, v, sz[v]);
    was[v] = true;
    vector<pair<int, int>> f;
    f.emplace_back(0, 0);                        
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (!was[u]) {
        Go(u, v, w, 1);  
        for (auto& p : vec) {
          f.push_back(p);
        }
        ans -= Calc(vec, k);
        vec.clear();
      }
    }
    ans += Calc(f, k);
    for (auto& e : g[v]) {
      int u = e.first;
      if (!was[u]) {
        Solve(u);
      }
    }
  };    
  Solve(0);
  cout << ans * 2 << '\n';                                            
  return 0;
}

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60