Submission #1159761

#	Time	Username	Problem	Language	Result	Execution time	Memory
1159761		yhkhoo	Janjetina (COCI21_janjetina)	C++20	50 / 110	35 ms	9284 KiB

Main

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

typedef pair<int, int> pii;
#define fi first
#define se second
#define mp make_pair
typedef pair<int, pii> pip;
#define pb push_back
#define eb emplace_back

const int MAXN = 100005;

int d[MAXN], m[MAXN], p[MAXN], n, k;
vector<pii> al[MAXN];
pip e[MAXN];

int main(){
	cin.tie(0); ios_base::sync_with_stdio(0);
	cin >> n >> k;
	for(int i=0; i<n-1; i++){
		cin >> e[i].se.fi >> e[i].se.se >> e[i].fi;
		if(e[i].se.fi > e[i].se.se) swap(e[i].se.fi, e[i].se.se);
		al[e[i].se.fi].eb(e[i].se.se, e[i].fi);
		al[e[i].se.se].eb(e[i].se.fi, e[i].fi);
	}
	sort(e, e+n-1);
	if(n <= 1000){
		long long ans = 0;
		for(int i=1; i<=n; i++){
			queue<int> q;
			memset(d, 0, (n+1)*sizeof(int));
			memset(m, 0, (n+1)*sizeof(int));
			memset(p, 0, (n+1)*sizeof(int));
			q.push(i);
			while(q.size()){
				int cn = q.front();
				q.pop();
				if(m[cn] - d[cn] >= k){
					// cerr << i << ' ' << cn << '\n';
					ans++;
				}
				for(auto j: al[cn]){
					auto [nn, w] = j;
					if(nn == p[cn]) continue;
					d[nn] = d[cn] + 1;
					m[nn] = max(m[cn], w);
					p[nn] = cn;
					q.push(nn);
				}
			}
		}
		cout << ans << '\n';
	}
	else{
		long long ans = 0;
		m[0] = 0;
		for(int i=1; i<n; i++){
			m[i] = m[i-1] + i;
		}
		for(int i=1; i<=n; i++){
			// p diddy
			p[i] = i;
			d[i] = i;
		}
		for(auto i = e; i != e+n-1; i++){
			auto [w, xy] = *i;
			auto [x, y] = xy;
			int a = x - p[x] + 1;
			int b = d[y] - y + 1;
			int l = w-k;
			//cerr << p[x] << ' ' << x << ' ' << y << ' ' << d[y] << '\n';
			//cerr << l << '\n';
			d[p[x]] = d[y];
			p[d[y]] = p[x];
			if(l < 1) continue;
			a = min(a, l);
			b = min(b, l);
			if(a > b) swap(a, b);
			//cerr << a << ' ' << b << '\n';
			ans += a*(min(b, l-a));
			ans += m[a] - m[min(a, l-b)];
			//cerr << ans << '\n';
		}
		ans *= 2;
		cout << ans << '\n';
	}
}

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