Submission #1221209

#	Time	Username	Problem	Language	Result	Execution time	Memory
1221209		tminh	Janjetina (COCI21_janjetina)	C++20	50 / 110	258 ms	35908 KiB

Main

#include "bits/stdc++.h"
using namespace std;

#define task ""
#define ll long long
#define endl '\n'
#define fi first
#define se second
#define vall(a) (a).begin(), (a).end()
#define sze(a) (int)a.size()
#define pii pair<int, int>
#define pll pair<ll, ll>


#define ep emplace_back
#define pb push_back
#define pf push_front


const ll mod = 1e9 + 7;
const int N = 2e5 + 5;
const ll oo = 1e18;

bool START;

int n, sz[N], k;
vector<pii> adj[N];
bool del[N];

int get_sz(int u, int p = 0) {
	sz[u] = 1;
	for (auto[v, w] : adj[u]) {
		if (v == p || del[v]) continue;
		sz[u] += get_sz(v, u);
	}
	return sz[u];
}

int get_cen(int sigma, int u, int p = 0) {
	for (auto[v, w] : adj[u]) {
		if (v == p || del[v]) continue;
		if (sz[v] > (sigma >> 1)) return get_cen(sigma, v, u);
	}
	return u;
}


vector<pii> arr[N], all;
void add(int u, int p, int depth, int Max, int cur) {
	arr[cur].pb(make_pair(Max, depth));
	all.pb(make_pair(Max, depth));
	for (auto[v, w] : adj[u]) {
		if (v == p || del[v]) continue;
		add(v, u, depth + 1, max(Max, w), cur);
	}
}


struct BIT {
	int bit[N];
	void add(int idx, int delta) {
		for (; idx <= n; idx += idx & -idx) bit[idx] += delta;
		return;
	}
	int get(int idx) {
		int res = 0;
		for (; idx > 0; idx -= idx & -idx) res += bit[idx];
		return res;
	}
	int get(int l, int r) {
		return get(r) - get(l - 1);
	}
} bit;

ll ans = 0;
void centroid(int u = 1) {
	u = get_cen(get_sz(u), u);
	del[u] = true;
	for (auto [v, w] : adj[u]) {
		if (del[v]) continue;
		add(v, u, 1, w, v);
		sort(vall(arr[v]), greater<pii>());
	}
	for (auto[v, w] : adj[u]) {
		if (del[v]) continue;
		for (auto[Max, depth] : arr[v]) {
			ans -= 1ll * bit.get(depth, n);
			if ((Max - depth - k) > 0 && (Max - depth - k) <= n) bit.add((Max - depth - k), 1);
		}
		for (auto[Max, depth] : arr[v]) if ((Max - depth - k) > 0 && (Max - depth - k) <= n) bit.add((Max - depth - k), -1);
		arr[v].clear();
	}
	
	sort(vall(all), greater<pii>());
	for (auto[Max, depth] : all) {
		ans += 1ll * bit.get(depth, n);
		if (Max - depth >= k) ++ans;
		if ((Max - depth - k) > 0 && (Max - depth - k) <= n) bit.add((Max - depth - k), 1);
	}
	for (auto[Max, depth] : all) if ((Max - depth - k) > 0 && (Max - depth - k) <= n) bit.add((Max - depth - k), -1);
	all.clear();
	
	
	for (auto [v, w] : adj[u]) {
		if (del[v]) continue;
		centroid(v);
	}
}

inline void solve() {
	centroid();
	cout << (ans << 1ll);
}


inline void input() {
	cin >> n >> k;
	for (int i = 1; i < n; ++i) {
		int u, v, c;
		cin >> u >> v >> c;
		adj[u].pb(make_pair(v, c));
		adj[v].pb(make_pair(u, c));
	}
	
    return solve();
}
bool END;

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

    input();
    
    
    cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << 's' << endl;
    cerr << "Memory: " << fabs ((&END - &START)) / 1048576.0 << "MB\n";
    return 0;
}

Compilation message (stderr)

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

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