Submission #366499

#	Time	Username	Problem	Language	Result	Execution time	Memory
366499		silverfish	Janjetina (COCI21_janjetina)	C++14	110 / 110	1025 ms	156100 KiB

Main

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

/*<DEBUG>*/
#define tem template <typename 
#define can_shift(_X_, ...) enable_if_t<sizeof test<_X_>(0) __VA_ARGS__ 8, debug&> operator<<(T i)
#define _op debug& operator<<
tem C > auto test(C *x) -> decltype(cerr << *x, 0LL);
tem C > char test(...);
tem C > struct itr{C begin, end; };
tem C > itr<C> get_range(C b, C e) { return itr<C>{b, e}; }
struct debug{
#ifdef _LOCAL
	~debug(){ cerr << endl; }
	tem T > can_shift(T, ==){ cerr << boolalpha << i; return *this; }
	tem T> can_shift(T, !=){ return *this << get_range(begin(i), end(i)); }
	tem T, typename U > _op (pair<T, U> i){ 
		return *this << "< " << i.first << " , " << i.second << " >"; }
	tem T> _op (itr<T> i){
		*this <<  "{ ";
		for(auto it = i.begin; it != i.end; it++){
			*this << " , " + (it==i.begin?2:0) << *it;
		}
		return *this << " }";
	}
#else
tem T> _op (const T&) { return *this; }
#endif 
};

string _ARR_(int* arr, int sz){
	string ret = "{ " + to_string(arr[0]); 
	for(int i = 1; i < sz; i++) ret += " , " +  to_string(arr[i]);
	ret += " }"; return ret;
}

#define exp(...) " [ " << #__VA_ARGS__ << " : " << (__VA_ARGS__) << " ]"
/*</DEBUG>*/

typedef long long ll;
typedef unsigned long long ull;
typedef unsigned int uint;
typedef pair<int, int> pii;
//mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

#define pb push_back
#define FAST ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0)
#define TC int __TC__; cin >> __TC__; while(__TC__--)
#define ar array

const int INF = 1e9 + 7, N = 1e5;
/*
 *
 * This solution could be better (maybe not faster, but consume less memory) by using only one fenwick tree
 *
 * Im not gonna implement that tho, my solution passes so its ok
 *
 */

int n, k, blocked[N], sz[N], d[N];
ll ans = 0;
vector<ar<int,2>> adj[N];
vector<ar<int,3>> ord[N];
int id=0, par[N], iter = 0, upper;
unordered_map<int, int> mx, fen[2*N];

void add(int i, int j){
	for(; i <= upper+1; i += i&-i){
		++fen[j][i];
	}
}

int get(int i, int j){
	int ret = 0;
	for(; i > 0; i -= i&-i){
		ret += fen[j][i];
	}
	return ret;
}
		
void dfs(int u, int p, int W){
	d[u] = d[p] + 1;
	upper = max(upper, d[u]);
	mx[u] = max(W, mx[p]);

	ord[iter].pb({mx[u], d[u], id});

	if(mx[u] - d[u] >= k) ++ans;

	for(auto [v ,w] : adj[u]){
		if(!blocked[v] && v != p){
			dfs(v, u, w);
		}
	}

}

void solve_root(int u){
	int all = id++;

	d[u] = 0;
	mx[u] = 0;
	upper = 0;
	for(auto [v, W] : adj[u]){
		if(blocked[v]) continue;
		dfs(v, u, W);
		id++;
	}

	sort(ord[iter].begin(), ord[iter].end());

	for(auto [w, l, id] : ord[iter]){
		if(w - l - k > 0){
			ans += get(min(w-l-k, upper+1), all) - get(min(w-l-k, upper+1), id);
		}
		add(l, all);
		add(l, id);
	}
	iter++;
	return;
}

void getsz(int u, int p){
	sz[u] = 1;
	par[u] = p;
	for(auto [v, w] : adj[u]){
		if(v != p && !blocked[v]){
			getsz(v, u);
			sz[u] += sz[v];
		}
	}
}

void decomp(int root){
	getsz(root, root);	
	if(sz[root] == 1) return;
	ar<int,2> centroid = {INF, -1};

	queue<int> q;
	q.push(root);

	while(q.size()){
		int u = q.front();
		q.pop();

		ar<int,2> mx_subtree = {sz[root] - sz[u], u};

		for(auto [v, w] : adj[u]){
			if(v != par[u] && !blocked[v]){
				q.push(v);
				mx_subtree = max(mx_subtree, {sz[v], u});
			}
		}
		centroid = min(centroid, mx_subtree);	
	}



	solve_root(centroid[1]);
//	debug() << exp(centroid[1]);

	blocked[centroid[1]] = 1;	

	for(auto [v, w] : adj[centroid[1]]){
		if(!blocked[v]) decomp(v);
	}

	return;
}

int main(void)
{
	FAST;
	cin >> n >> k;
	for(int i = 1; i < n; ++i){
		int u, v, w; cin >> u >> v >> w;
		--u; --v;
		adj[u].pb({v,w});
		adj[v].pb({u,w});
	}

	decomp(0);

	cout << ans*2 << '\n';

	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void dfs(int, int, int)':
Main.cpp:90:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
   90 |  for(auto [v ,w] : adj[u]){
      |           ^
Main.cpp: In function 'void solve_root(int)':
Main.cpp:104:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  104 |  for(auto [v, W] : adj[u]){
      |           ^
Main.cpp:112:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  112 |  for(auto [w, l, id] : ord[iter]){
      |           ^
Main.cpp: In function 'void getsz(int, int)':
Main.cpp:126:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  126 |  for(auto [v, w] : adj[u]){
      |           ^
Main.cpp: In function 'void decomp(int)':
Main.cpp:148:12: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  148 |   for(auto [v, w] : adj[u]){
      |            ^
Main.cpp:164:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  164 |  for(auto [v, w] : adj[centroid[1]]){
      |           ^

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