verify/jump_on_tree-level_ancestor.test.cpp

View this file on GitHub · Last update: 2026-07-25 01:16:11-07:00

Problem: https://judge.yosupo.jp/problem/jump_on_tree

Depends on

• level_ancestor.hpp
• yc.hpp

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/jump_on_tree

#include <bits/stdc++.h>
#include <cassert>

#include "level_ancestor.hpp"

int main() {
	std::ios_base::sync_with_stdio(false); std::cin.tie(nullptr);

	int N, Q; std::cin >> N >> Q;
	std::vector<std::vector<int>> adj(N);
	for (int e = 0; e < N-1; e++) {
		int u, v; std::cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	std::vector<int> par(N);
	std::vector<int> depth(N);
	[&](this auto&& self, int cur, int prv, int d) -> void {
		par[cur] = prv;
		depth[cur] = d;
		for (int nxt : adj[cur]) {
			if (nxt == prv) continue;
			self(nxt, cur, d+1);
		}
	}(0, -1, 0);
	ecnerwala::level_ancestor la(par);


	for (int q = 0; q < Q; q++) {
		int s, t, i; std::cin >> s >> t >> i;
		std::cout << [&]() -> int {
			int l = la.lca(s, t);
			if (i <= depth[s] - depth[l]) {
				return la.get_ancestor(s, i);
			}
			i -= depth[s] - depth[l];
			if (i <= depth[t] - depth[l]) {
				return la.get_ancestor(t, depth[t] - depth[l] - i);
			}
			return -1;
		}() << '\n';
	}

	return 0;
}

#include <bits/stdc++.h>
#line 1 "verify/jump_on_tree-level_ancestor.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/jump_on_tree

#line 5 "verify/jump_on_tree-level_ancestor.test.cpp"

#line 2 "src/level_ancestor.hpp"

#line 6 "src/level_ancestor.hpp"

#line 2 "src/yc.hpp"

#line 5 "src/yc.hpp"

namespace std {

template<class Fun>
class y_combinator_result {
	Fun fun_;
public:
	template<class T>
	explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}

	template<class ...Args>
	decltype(auto) operator()(Args &&...args) {
		return fun_(std::ref(*this), std::forward<Args>(args)...);
	}
};

template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
	return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

} // namespace std
#line 8 "src/level_ancestor.hpp"

namespace ecnerwala {

using std::swap;

struct level_ancestor {
	int N;
	std::vector<int> preorder;
	std::vector<int> idx;
	std::vector<std::pair<int, int>> heavyPar; // heavy parent, distance
	level_ancestor() : N(0) {}

	level_ancestor(const std::vector<int>& par) : N(int(par.size())), preorder(N), idx(N), heavyPar(N) {
		std::vector<std::vector<int>> ch(N);
		for (int i = 0; i < N; i++) {
			if (par[i] != -1) ch[par[i]].push_back(i);
		}
		std::vector<int> sz(N);
		int nxt_idx = 0;
		for (int i = 0; i < N; i++) {
			if (par[i] == -1) {
				std::y_combinator([&](auto self, int cur) -> void {
					sz[cur] = 1;
					for (int nxt : ch[cur]) {
						self(nxt);
						sz[cur] += sz[nxt];
					}
					if (!ch[cur].empty()) {
						auto mit = std::max_element(ch[cur].begin(), ch[cur].end(), [&](int a, int b) { return sz[a] < sz[b]; });
						swap(*ch[cur].begin(), *mit);
					}
				})(i);
				std::y_combinator([&](auto self, int cur, int isRoot = true) -> void {
					preorder[idx[cur] = nxt_idx++] = cur;
					if (isRoot) {
						heavyPar[idx[cur]] = {par[cur] == -1 ? -1 : idx[par[cur]], 1};
					} else {
						assert(idx[par[cur]] == idx[cur]-1);
						heavyPar[idx[cur]] = heavyPar[idx[cur]-1];
						heavyPar[idx[cur]].second++;
					}
					bool chRoot = false;
					for (int nxt : ch[cur]) {
						self(nxt, chRoot);
						chRoot = true;
					}
				})(i);
			}
		}
	}

	int get_ancestor(int a, int k) const {
		assert(k >= 0);
		a = idx[a];
		while (a != -1 && k) {
			if (k >= heavyPar[a].second) {
				k -= heavyPar[a].second;
				assert(heavyPar[a].first <= a - heavyPar[a].second);
				a = heavyPar[a].first;
			} else {
				a -= k;
				k = 0;
			}
		}
		if (a == -1) return -1;
		else return preorder[a];
	}

	int lca(int a, int b) const {
		a = idx[a], b = idx[b];
		while (true) {
			if (a > b) swap(a, b);
			assert(a <= b);
			if (a > b - heavyPar[b].second) {
				return preorder[a];
			}
			b = heavyPar[b].first;
			if (b == -1) return -1;
		}
	}

	int dist(int a, int b) const {
		a = idx[a], b = idx[b];
		int res = 0;
		while (true) {
			if (a > b) swap(a, b);
			assert(a <= b);
			if (a > b - heavyPar[b].second) {
				res += b - a;
				break;
			}
			res += heavyPar[b].second;
			b = heavyPar[b].first;
			if (b == -1) return -1;
		}
		return res;
	}
};

} // namespace ecnerwala
#line 7 "verify/jump_on_tree-level_ancestor.test.cpp"

int main() {
	std::ios_base::sync_with_stdio(false); std::cin.tie(nullptr);

	int N, Q; std::cin >> N >> Q;
	std::vector<std::vector<int>> adj(N);
	for (int e = 0; e < N-1; e++) {
		int u, v; std::cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	std::vector<int> par(N);
	std::vector<int> depth(N);
	[&](this auto&& self, int cur, int prv, int d) -> void {
		par[cur] = prv;
		depth[cur] = d;
		for (int nxt : adj[cur]) {
			if (nxt == prv) continue;
			self(nxt, cur, d+1);
		}
	}(0, -1, 0);
	ecnerwala::level_ancestor la(par);


	for (int q = 0; q < Q; q++) {
		int s, t, i; std::cin >> s >> t >> i;
		std::cout << [&]() -> int {
			int l = la.lca(s, t);
			if (i <= depth[s] - depth[l]) {
				return la.get_ancestor(s, i);
			}
			i -= depth[s] - depth[l];
			if (i <= depth[t] - depth[l]) {
				return la.get_ancestor(t, depth[t] - depth[l] - i);
			}
			return -1;
		}() << '\n';
	}

	return 0;
}

// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
// src/yc.hpp
namespace std{
template<class Fun>
class y_combinator_result{
Fun fun_;
public:
template<class T>
explicit y_combinator_result(T&&fun):fun_(std::forward<T>(fun)){}
template<class...Args>
decltype(auto)operator()(Args&&...args){
return fun_(std::ref(*this),std::forward<Args>(args)...);
}
};
template<class Fun>
decltype(auto)y_combinator(Fun&&fun){
return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}
}
// src/level_ancestor.hpp
namespace ecnerwala{
using std::swap;
struct level_ancestor{
int N;
std::vector<int>preorder;
std::vector<int>idx;
std::vector<std::pair<int,int>>heavyPar;
level_ancestor():N(0){}
level_ancestor(const std::vector<int>&par):N(int(par.size())),preorder(N),idx(N),heavyPar(N){
std::vector<std::vector<int>>ch(N);
for(int i=0;i<N;i++){
if(par[i]!=-1)ch[par[i]].push_back(i);
}
std::vector<int>sz(N);
int nxt_idx=0;
for(int i=0;i<N;i++){
if(par[i]==-1){
std::y_combinator([&](auto self,int cur)->void{
sz[cur]=1;
for(int nxt:ch[cur]){
self(nxt);
sz[cur]+=sz[nxt];
}
if(!ch[cur].empty()){
auto mit=std::max_element(ch[cur].begin(),ch[cur].end(),[&](int a,int b){return sz[a]<sz[b];});
swap(*ch[cur].begin(),*mit);
}
})(i);
std::y_combinator([&](auto self,int cur,int isRoot=true)->void{
preorder[idx[cur]=nxt_idx++]=cur;
if(isRoot){
heavyPar[idx[cur]]={par[cur]==-1?-1:idx[par[cur]],1};
}else{
assert(idx[par[cur]]==idx[cur]-1);
heavyPar[idx[cur]]=heavyPar[idx[cur]-1];
heavyPar[idx[cur]].second++;
}
bool chRoot=false;
for(int nxt:ch[cur]){
self(nxt,chRoot);
chRoot=true;
}
})(i);
}
}
}
int get_ancestor(int a,int k)const{
assert(k>=0);
a=idx[a];
while(a!=-1&&k){
if(k>=heavyPar[a].second){
k-=heavyPar[a].second;
assert(heavyPar[a].first<=a-heavyPar[a].second);
a=heavyPar[a].first;
}else{
a-=k;
k=0;
}
}
if(a==-1)return-1;
else return preorder[a];
}
int lca(int a,int b)const{
a=idx[a],b=idx[b];
while(true){
if(a>b)swap(a,b);
assert(a<=b);
if(a>b-heavyPar[b].second){
return preorder[a];
}
b=heavyPar[b].first;
if(b==-1)return-1;
}
}
int dist(int a,int b)const{
a=idx[a],b=idx[b];
int res=0;
while(true){
if(a>b)swap(a,b);
assert(a<=b);
if(a>b-heavyPar[b].second){
res+=b-a;
break;
}
res+=heavyPar[b].second;
b=heavyPar[b].first;
if(b==-1)return-1;
}
return res;
}
};
}
// verify/jump_on_tree-level_ancestor.test.cpp
int main(){
std::ios_base::sync_with_stdio(false);std::cin.tie(nullptr);
int N,Q;std::cin>>N>>Q;
std::vector<std::vector<int>>adj(N);
for(int e=0;e<N-1;e++){
int u,v;std::cin>>u>>v;
adj[u].push_back(v);
adj[v].push_back(u);
}
std::vector<int>par(N);
std::vector<int>depth(N);
[&](this auto&&self,int cur,int prv,int d)->void{
par[cur]=prv;
depth[cur]=d;
for(int nxt:adj[cur]){
if(nxt==prv)continue;
self(nxt,cur,d+1);
}
}(0,-1,0);
ecnerwala::level_ancestor la(par);
for(int q=0;q<Q;q++){
int s,t,i;std::cin>>s>>t>>i;
std::cout<<[&]()->int{
int l=la.lca(s,t);
if(i<=depth[s]-depth[l]){
return la.get_ancestor(s,i);
}
i-=depth[s]-depth[l];
if(i<=depth[t]-depth[l]){
return la.get_ancestor(t,depth[t]-depth[l]-i);
}
return-1;
}()<<'\n';
}
return 0;
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on

Test cases

Env	Name	Status	Elapsed	Memory
g++-sanitizer	almost_line_00	AC	4625 ms	358 MB
g++-sanitizer	almost_line_01	AC	4555 ms	276 MB
g++-sanitizer	almost_line_02	AC	4554 ms	267 MB
g++-sanitizer	almost_line_03	AC	4550 ms	257 MB
g++-sanitizer	almost_uni_00	AC	602 ms	93 MB
g++-sanitizer	almost_uni_01	AC	559 ms	93 MB
g++-sanitizer	example_00	AC	13 ms	9 MB
g++-sanitizer	line_00	AC	4509 ms	659 MB
g++-sanitizer	max_random_00	AC	1098 ms	118 MB
g++-sanitizer	max_random_01	AC	1104 ms	118 MB
g++-sanitizer	random_00	AC	877 ms	95 MB
g++-sanitizer	random_01	AC	1041 ms	111 MB
g++-sanitizer	random_02	AC	230 ms	26 MB
g++-sanitizer	random_03	AC	971 ms	104 MB
g++-sanitizer	random_04	AC	601 ms	72 MB
g++-sanitizer	small_random_00	AC	114 ms	12 MB
g++-sanitizer	small_random_01	AC	109 ms	12 MB
g++-sanitizer	small_random_02	AC	117 ms	12 MB
g++-sanitizer	small_random_03	AC	117 ms	12 MB
g++-sanitizer	small_random_04	AC	107 ms	12 MB
g++-sanitizer	uni_00	AC	514 ms	89 MB
g++	almost_line_00	AC	488 ms	76 MB
g++	almost_line_01	AC	478 ms	73 MB
g++	almost_line_02	AC	453 ms	72 MB
g++	almost_line_03	AC	485 ms	72 MB
g++	almost_uni_00	AC	260 ms	61 MB
g++	almost_uni_01	AC	251 ms	61 MB
g++	example_00	AC	2 ms	3 MB
g++	line_00	AC	370 ms	95 MB
g++	max_random_00	AC	404 ms	65 MB
g++	max_random_01	AC	411 ms	65 MB
g++	random_00	AC	315 ms	51 MB
g++	random_01	AC	376 ms	60 MB
g++	random_02	AC	124 ms	10 MB
g++	random_03	AC	363 ms	56 MB
g++	random_04	AC	249 ms	37 MB
g++	small_random_00	AC	70 ms	4 MB
g++	small_random_01	AC	68 ms	4 MB
g++	small_random_02	AC	69 ms	4 MB
g++	small_random_03	AC	73 ms	4 MB
g++	small_random_04	AC	73 ms	4 MB
g++	uni_00	AC	234 ms	60 MB

Back to top page