`#include "level_ancestor.hpp"`

View this file on GitHub · Last update: 2026-06-07 11:59:34-07:00

Depends on

yc.hpp

Verified with

verify/jump_on_tree-level_ancestor.test.cpp

Code

	Coverage	Exec / Excl / Total
Lines	78.1%	50 / 0 / 64
Functions	62.5%	5 / 0 / 8
Branches	40.4%	44 / 0 / 109
Full report

#pragma once

#include <algorithm>
#include <vector>
#include <cassert>

#include "yc.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

#include <algorithm>
#include <vector>
#include <cassert>
#include <functional>
#include <utility>
#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

// 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>
#include <cassert>
// 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;
}
};
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on

#pragma once

#include <algorithm>
#include <vector>
#include <cassert>

#include "yc.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