Submission #366951

#	Time	Username	Problem	Language	Result	Execution time	Memory
366951		alrad	Janjetina (COCI21_janjetina)	C++17	110 / 110	597 ms	16232 KiB

Main

#include <bits/stdc++.h>

using namespace std;

using ld = long double;
using ull = unsigned long long;

/*
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("-O3")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")
*/

template <class T> inline T gcd(T a , T b) { return !a ? b : gcd(b % a , a); }
template <class T> inline T lcm(T a , T b) {return (a * b) / gcd(a , b) ; }

mt19937 rnd(time(0));

#define all(x) x.begin(), x.end()
#define debug(x) { cerr << #x << " = " << x << endl; }

const int N = 1e5 + 2;

int n, k, sub[N], depth[N];
long long fen[N], ans = 0ll;
vector<bool> used(N, false);

vector<pair<int, int>> cnt;
vector<vector<pair<int, int>>> g(N, vector<pair<int, int>>());

void calc(int v, int par = -1) {
  sub[v] = 1;
  if (par != -1) {
    depth[v] = depth[par] + 1;
  } else {
    depth[v] = 0;
  }
  for (auto edge : g[v]) {
    int to = edge.first;
    if (to != par && !used[to]) {
      calc(to, v);
      sub[v] += sub[to];
    }
  }
  return;
}

int getCentroid(int v, int par, int subSize) {
  for (auto edge : g[v]) {
    int to = edge.first;
    if (sub[to] > (subSize / 2) && !used[to] && to != par) {
      return getCentroid(to, v, subSize);
    }
  }
  return v;
}

void dfs1(int v, int par, int maxEdge) {
  cnt.push_back({v, maxEdge});
  for (auto edge : g[v]) {
    int to = edge.first, wei = edge.second;
    if (!used[to] && to != par) {
      dfs1(to, v, max(maxEdge, wei));
    }
  }
  return;
}

bool comp(const pair<int, int> &p1, const pair<int, int> &p2) {
  if (p1.second != p2.second) {
    return p1.second < p2.second;
  }
  return p1.first < p2.first;
}

long long query(int R) {
  long long sum = 0ll;
  R++;
  for (; R >= 0; R = (R & (R + 1)) - 1) {
    sum += fen[R];
  }
  return sum;
}

void upd(int i, long long value) {
  i++;
  for (; i < N - 1; i = (i | (i + 1))) {
    fen[i] += value;
  }
  return;
}

void decompose(int v) {
  calc(v);
  used[v] = true;
  dfs1(v, -1, 0);
  sort(all(cnt), comp);
  for (auto temp : cnt) {
    int ver = temp.first, wei = temp.second;
    ans += query(wei - k - depth[ver]);
    upd(depth[ver], +1);
  }
  for (auto temp : cnt) {
    int ver = temp.first;
    upd(depth[ver], -1);
  }
  cnt.clear();
  for (auto edge : g[v]) {
    int to = edge.first;
    if (!used[to]) {
      dfs1(to, v, edge.second);
      sort(all(cnt), comp);
      for (auto temp : cnt) {
        int ver = temp.first, wei = temp.second;
        ans -= query(wei - k - depth[ver]);
        upd(depth[ver], +1);
      }
      for (auto temp : cnt) {
        int ver = temp.first;
        upd(depth[ver], -1);
      }
      cnt.clear();
    }
  }
  for (auto edge : g[v]) {
    int to = edge.first;
    if (!used[to]) {
      decompose(getCentroid(to, v, sub[to]));
    }
  }
  return;
}

void solve() {
  cin >> n >> k;
  for (int i = 1; i < n; i++) {
    int u, v, wei;
    cin >> u >> v >> wei;
    g[u].push_back({v, wei});
    g[v].push_back({u, wei});
  }
  decompose(getCentroid(1, -1, n));
  cout << 2ll * ans << '\n';
  return;
}

signed main() {
  ios_base :: sync_with_stdio(0);
  cin.tie(0) , cout.tie(0);
  int t = 1;
  while (t-- > 0) {
    solve();
  }
  return 0;
}

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