Submission #877280

#	Time	Username	Problem	Language	Result	Execution time	Memory
877280		mgl_diamond	Janjetina (COCI21_janjetina)	C++17	110 / 110	149 ms	20032 KiB

Main

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ii = pair<int, int>;
 
#define foru(i, l, r) for(int i=(l); i<=(r); ++i)
#define ford(i, l, r) for(int i=(l); i>=(r); --i)
#define fore(x, v) for(auto &x : v)
#define all(x) (x).begin(), (x).end()
#define sz(x) (int)(x).size()
#define fi first
#define se second
#define file "input"

template<class T> bool minimize(T &a, T b) { if (a > b) { a = b; return 1; } return 0; }
template<class T> bool maximize(T &a, T b) { if (a < b) { a = b; return 1; } return 0; }

void setIO() {
  ios::sync_with_stdio(0);
  cin.tie(0); cout.tie(0);
  if (fopen(file".inp", "r")) {
    freopen(file".inp", "r", stdin);
    freopen(file".out", "w", stdout);
  }
}

const int N = 1e5+5;

int n, k;
ll ans;
vector<ii> adj[N];

// centroid
int cent, tree_size, max_dep;
  
struct fenwick {
  int tmp;
  vector<int> bit;

  void init(int n) {
    tmp = n;
    bit = vector<int>(n+2, 0);
  }

  void update(int i, int v) {
    assert(i <= tmp);
    for(++i; i<=tmp+1; i+=i&-i) bit[i] += v;
  }

  int query(int i) {
    i = min(tmp, i);
    int ans = 0;
    for(++i; i>=1; i-=i&-i) ans += bit[i];
    return ans;
  }
} branch[N], al;

int high[N], sub[N], who[N];
bool process[N];

int get_subsize(int u, int p) {
  sub[u] = 1;
  fore(x, adj[u]) {
    int v = x.fi;
    if (v == p || process[v]) continue;
    sub[u] += get_subsize(v, u);
  }
  return sub[u];
}

int get_centroid(int u, int p) {
  fore(x, adj[u]) {
    int v = x.fi;
    if (v == p || process[v]) continue;
    if (sub[v] * 2 > tree_size) return get_centroid(v, u);
  }
  return u;
}

vector<ii> cands;

int dfs_prep(int u, int p, int mark, int mx = 0) {
  if (mx - high[u] >= k) ++ans;
  if (u != cent) {
    int ans = 0;
    who[u] = mark;
    fore(x, adj[u]) {
      int v = x.fi, w = x.se;
      if (v == p || process[v]) continue;
      cands.emplace_back(max(mx, w), v);
      high[v] = high[u] + 1;
      maximize(ans, dfs_prep(v, u, mark, max(mx, w)) + 1);
    }
    return ans;
  }
  else {
    int ans = 0;
    fore(x, adj[u]) {
      int v = x.fi, w = x.se;
      if (v == p || process[v]) continue;
      cands.emplace_back(max(mx, w), v);
      high[v] = high[u] + 1;
      int tmp = dfs_prep(v, u, mark, max(mx, w));
      maximize(ans, tmp + 1);
      branch[mark].init(tmp+1);
      ++mark;
    }
    return ans;
  }
}

void solve(int u) {
  tree_size = get_subsize(u, 0);
  cent = get_centroid(u, 0);
  cands.clear();
  high[cent] = 0;
  max_dep = dfs_prep(cent, cent, 0);
  al.init(max_dep);

  // cout << ans << "\n";
  sort(all(cands));
  fore(x, cands) {
    int j = x.fi - high[x.se] - k;
    if (j >= 0) ans += al.query(j) - branch[who[x.se]].query(j);
    // if (x.se == 2) cout << al.query(j) << "\n";
    // cout << x.fi - high[x.se] - k << "\n";
    al.update(high[x.se], 1);
    branch[who[x.se]].update(high[x.se], 1);
  } 

  // cout << cent << " " << ans << "\n";
  process[cent] = 1;
  fore(x, adj[cent]) {
    int v = x.fi, w = x.se;
    if (process[v]) continue;
    solve(v);
  }
}

int main() {
  setIO();

  cin >> n >> k;
  foru(i, 1, n-1) {
    int u, v, w;
    cin >> u >> v >> w; 
    adj[u].emplace_back(v, w);
    adj[v].emplace_back(u, w);
  }

  solve(1);  
  cout << ans*2;
}

/*
5 2
1 2 5
2 3 2
1 4 3
3 5 4
*/

Compilation message (stderr)

Main.cpp: In function 'void solve(int)':
Main.cpp:134:19: warning: unused variable 'w' [-Wunused-variable]
  134 |     int v = x.fi, w = x.se;
      |                   ^
Main.cpp: In function 'void setIO()':
Main.cpp:22:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   22 |     freopen(file".inp", "r", stdin);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:23:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   23 |     freopen(file".out", "w", stdout);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~

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