Submission #1191915

#	Time	Username	Problem	Language	Result	Execution time	Memory
1191915		SmuggingSpun	Janjetina (COCI21_janjetina)	C++20	110 / 110	281 ms	17032 KiB

Main

#include<bits/stdc++.h>
#define taskname "D"
using namespace std;
typedef long long ll;
const int lim = 1e5 + 5;
int n, k, u[lim], v[lim], w[lim];
vector<int>g[lim];
namespace sub1{
	int ans = 0;
	void dfs(int s, int h = 0, int max_w = 0, int p = -1){
		if(max_w - h >= k){
			ans++;
		}
		for(int& i : g[s]){
			int d = u[i] ^ v[i] ^ s;
			if(d != p){
				dfs(d, h + 1, max(max_w, w[i]), s);
			}
		}
	}
	void solve(){
		for(int i = 1; i <= n; i++){
			dfs(i);
		}
		cout << ans;
	}
}
namespace sub23{
	int cnt_child[lim], bit[lim];
	bitset<lim>vis;
	void update(int p, int x){
		for(p++; p < lim; p += p & -p){
			bit[p] += x;
		}
	}
	int get(int p){
		int ans = 0;
		for(p++; p > 0; p -= p & -p){
			ans += bit[p];
		}
		return ans;
	}
	void dfs_1(int s, int p = -1){
		cnt_child[s] = 1;
		for(int& i : g[s]){
			int d = u[i] ^ v[i] ^ s;
			if(d != p && !vis.test(d)){
				dfs_1(d, s);
				cnt_child[s] += cnt_child[d];
			}
		}
	}
	int get_centroid(int s){
		dfs_1(s);
		int N = cnt_child[s], p = -1;
		while(true){
			bool flag = true;
			for(int& i : g[s]){
				int d = u[i] ^ v[i] ^ s;
				if(!vis.test(d) && d != p && cnt_child[d] > (N >> 1)){
					p = s;
					s = d;
					flag = false;
					break;
				}
			}
			if(flag){
				break;
			}
		}
		return s;
	}
	vector<pair<int, int>>value;
	void dfs_2(int s, int h = 0, int max_w = 0, int p = -1){
		value.emplace_back(max_w, h);
		for(int& i : g[s]){
			int d = u[i] ^ v[i] ^ s;
			if(!vis.test(d) && d != p){
				dfs_2(d, h + 1, max(max_w, w[i]), s);
			}
		}
	}
	ll ans = 0;
	void solve(bool positive){
		sort(value.begin(), value.end());
		for(int i = 0; i < value.size(); i++){
			int add = get(value[i].first - value[i].second - k);
			if(positive){
				ans += add;
			}
			else{
				ans -= add;
			}
			update(value[i].second, 1);
		}
		for(int i = 0; i < value.size(); i++){
			update(value[i].second, -1);
		}
		value.clear();
	}
	void centroid_decomposition(int s){
		int centroid = get_centroid(s);
		dfs_2(centroid);
		solve(true);
		vis.set(centroid);
		for(int& i : g[centroid]){
			int d = u[i] ^ v[i] ^ centroid;
			if(!vis.test(d)){
				dfs_2(d, 1, w[i], centroid);
				solve(false);
			}
		}
		for(int& i : g[centroid]){
			int d = u[i] ^ v[i] ^ centroid;
			if(!vis.test(d)){
				centroid_decomposition(d);
			}
		}
	}
	void solve(){
		memset(bit, 0, sizeof(bit));
		vis.reset();
		centroid_decomposition(1);
		cout << (ans << 1LL);
	}
}
int main(){
	ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	if(fopen(taskname".inp", "r")){
		freopen(taskname".inp", "r", stdin);
	}
	cin >> n >> k;
	for(int i = 1; i < n; i++){
		cin >> u[i] >> v[i] >> w[i];
		g[u[i]].emplace_back(i);
		g[v[i]].emplace_back(i);
	}
	if(n <= 1000){
		sub1::solve();
	}
	else{
		sub23::solve();
	}
}

Compilation message (stderr)

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

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