verify/spoj/spoj-QTREE2.test.cpp

View this file on GitHub
Last update: 2022-07-21 16:12:33-04:00
Problem: https://www.spoj.com/problems/QTREE2/

Depends on

Code

#define PROBLEM "https://www.spoj.com/problems/QTREE2/"

#include "../../library/contest/template-short.hpp"
#include "../../library/graphs/lca-jump-distance.hpp"

void solve_case(int tc = 0) {
	int n;
	cin >> n;
	LCAJumpDistance<long long> lca;
	lca.init(n);
	f0r(i, n - 1) {
		int u, v, w;
		cin >> u >> v >> w;
		--u, --v;
		lca.ae(u, v, w);
	}
	lca.gen(0);
	while (true) {
		string s;
		cin >> s;
		if (s == "DIST") {
			int u, v;
			cin >> u >> v;
			--u, --v;
			cout << (lca.distance(u, v)) << '\n';
		} else if (s == "KTH") {
			int u, v, w;
			cin >> u >> v >> w;
			--u, --v, --w;
			int up = lca.lca(u, v);						
			int d1 = lca.depth[u] - lca.depth[up];
			int d2 = lca.depth[v] - lca.depth[up];
			if (d1 >= w) {
				cout << (lca.jump(u, w) + 1) << '\n';
			} else {
				w -= d1;
				w = d2 - w;
				cout << (lca.jump(v, w) + 1) << '\n';
			}
		} else {
			break;
		}
	}
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int tt = 1;
	cin >> tt;
	f1r(tc, 1, tt + 1) {
		solve_case(tc);
	}
	return 0;
}

#define PROBLEM "https://www.spoj.com/problems/QTREE2/"


#include <bits/stdc++.h>

using namespace std;

#define f1r(i, a, b) for (int i = (a); i < (b); ++i)
#define f0r(i, a) f1r(i, 0, a)
#define each(t, a) for (auto& t : a)

#define mp make_pair
#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define sz(x) (int)(x).size()
#define all(x) begin(x), end(x)

typedef long long ll;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef vector<pi> vpi;
typedef vector<pl> vpl;

template <class T> bool ckmin(T& a, const T& b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T& a, const T& b) { return a < b ? a = b, 1 : 0; }

template<class T> struct LCAJumpDistance {
	int n;
	std::vector<std::vector<int>> par;
	std::vector<std::vector<std::pair<int, T>>> adj;
	std::vector<int> depth;
	std::vector<T> depth_dist;

	void init(int _n) {
		n = _n;
		int d = 1;
		while ((1 << d) < n) d++;
		par.assign(d, std::vector<int>(n));
		adj.resize(n);
		depth.resize(n);
		depth_dist.resize(n);
	}

	void ae(int x, int y, T c = 1) {
		adj[x].emplace_back(y, c);
		adj[y].emplace_back(x, c);
	}

	void gen(int root = 0) {
		par[0][root] = root;
		dfs(root);
	}

	void dfs(int src = 0) {
		for (int i = 1; i < (int)par.size(); i++) {
			par[i][src] = par[i - 1][par[i - 1][src]];
		}
		for (auto nxt: adj[src]) {
			if (nxt.first == par[0][src]) continue;
			depth_dist[nxt.first] = depth_dist[par[0][nxt.first] = src] + nxt.second;
			depth[nxt.first] = depth[par[0][nxt.first] = src] + 1;
			dfs(nxt.first);
		}
	}

	int jump(int x, int d) {
		for (int i = 0; i < (int)par.size(); i++) {
			if ((d >> i) & 1) {
				x = par[i][x];
			}
		}
		return x;
	}

	int lca(int x, int y) {
		if (depth[x] < depth[y]) std::swap(x, y);
		x = jump(x, depth[x] - depth[y]);
		if (x == y) return x;
		for (int i = (int)par.size() - 1; i >= 0; i--) {
			int nx = par[i][x];
			int ny = par[i][y];
			if (nx != ny) x = nx, y = ny;
		}
		return par[0][x];
	}

	T distance(int x, int y) {
		int l = lca(x, y);
		return depth_dist[x] + depth_dist[y] - 2 * depth_dist[l];
	}
};

void solve_case(int tc = 0) {
	int n;
	cin >> n;
	LCAJumpDistance<long long> lca;
	lca.init(n);
	f0r(i, n - 1) {
		int u, v, w;
		cin >> u >> v >> w;
		--u, --v;
		lca.ae(u, v, w);
	}
	lca.gen(0);
	while (true) {
		string s;
		cin >> s;
		if (s == "DIST") {
			int u, v;
			cin >> u >> v;
			--u, --v;
			cout << (lca.distance(u, v)) << '\n';
		} else if (s == "KTH") {
			int u, v, w;
			cin >> u >> v >> w;
			--u, --v, --w;
			int up = lca.lca(u, v);						
			int d1 = lca.depth[u] - lca.depth[up];
			int d2 = lca.depth[v] - lca.depth[up];
			if (d1 >= w) {
				cout << (lca.jump(u, w) + 1) << '\n';
			} else {
				w -= d1;
				w = d2 - w;
				cout << (lca.jump(v, w) + 1) << '\n';
			}
		} else {
			break;
		}
	}
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int tt = 1;
	cin >> tt;
	f1r(tc, 1, tt + 1) {
		solve_case(tc);
	}
	return 0;
}