Submission #1286229

#	Time	Username	Problem	Language	Result	Execution time	Memory
1286229		lmquan	Janjetina (COCI21_janjetina)	C++20	110 / 110	258 ms	16392 KiB

Main

#define taskname ""
#include <bits/stdc++.h>
using namespace std;
const int kMaxNW = 100005;

struct FenwickTree {
  int x;
  vector<int> bit;

  FenwickTree(int _x) : x(_x) {
    bit.resize(x + 2, 0);
  }

  void Update(int p, int d) {
    p++;
    for (int i = p; i <= x; i += i & (-i)) {
      bit[i] += d;
    }
  }

  int PrefixSum(int p) {
    p++;
    int s = 0;
    for (int i = p; i > 0; i -= i & (-i)) {
      s += bit[i];
    }
    return s;
  }
};

int n, k, sz[kMaxNW];
long long cnt;
bool deleted[kMaxNW];
vector<pair<int, int>> adj[kMaxNW];
FenwickTree a(kMaxNW);

int DFS1(int u, int p) {
  sz[u] = 1;
  for (auto i : adj[u]) {
    int v = i.first;
    if (v == p || deleted[v]) {
      continue;
    }
    DFS1(v, u);
    sz[u] += sz[v];
  }
  return sz[u];
}

int FindCentroid(int u, int p, int m) {
  for (auto i : adj[u]) {
    int v = i.first;
    if (v == p || deleted[v]) {
      continue;
    }
    if (sz[v] > m / 2) {
      return FindCentroid(v, u, m);
    }
  }
  return u;
}

void DFS2(int u, int p, int depth, int max_weight, vector<pair<int, int>>& S) {
  S.push_back({depth, max_weight});
  for (auto i : adj[u]) {
    int v = i.first, w = i.second;
    if (v == p || deleted[v]) {
      continue;
    }
    DFS2(v, u, depth + 1, max(max_weight, w), S);
  }
}

void CentroidDecomposition(int u) {
  int g = DFS1(u, -1), c = FindCentroid(u, -1, g);
  deleted[c] = true;

  vector<pair<int, int>> S;
  DFS2(c, -1, 0, 0, S);
  sort(S.begin(), S.end(), [](pair<int, int> x, pair<int, int> y) {
    return x.second < y.second;
  });
  for (int i = 0; i < S.size(); i++) {
    if (S[i].second - S[i].first - k >= 0) {
      cnt += a.PrefixSum(S[i].second - S[i].first - k);
    }
    a.Update(S[i].first, 1);
  }
  for (int i = 0; i < S.size(); i++) {
    a.Update(S[i].first, -1);
  }

  for (auto i : adj[c]) {
    int v = i.first, w = i.second;
    if (deleted[v]) {
      continue;
    }
    vector<pair<int, int>> M;
    DFS2(v, c, 1, w, M);
    sort(M.begin(), M.end(), [](pair<int, int> x, pair<int, int> y) {
      return x.second < y.second;
    });
    for (int i = 0; i < M.size(); i++) {
      if (M[i].second - M[i].first - k >= 0) {
        cnt -= a.PrefixSum(M[i].second - M[i].first - k);
      }
      a.Update(M[i].first, 1);
    }
    for (int i = 0; i < M.size(); i++) {
      a.Update(M[i].first, -1);
    }
  }

  for (auto i : adj[c]) {
    int v = i.first;
    if (deleted[v]) {
      continue;
    }
    CentroidDecomposition(v);
  }
}

int main() {
  if (fopen(taskname".inp", "r")) {
    freopen(taskname".inp", "r", stdin);
    freopen(taskname".out", "w", stdout);
  }
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

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

  int root = 1;
  CentroidDecomposition(root);

  cout << 2 * cnt;

  return 0;
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:125:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  125 |     freopen(taskname".inp", "r", stdin);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:126:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  126 |     freopen(taskname".out", "w", stdout);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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