Test/EulerTour.path.test.cpp

• View this file on GitHub
• Last update: 2022-05-15 21:01:40+09:00
• Problem: https://yukicoder.me/problems/no/1094

Depends on

• Binary Indexed Tree (DataStructure/BinaryIndexedTree.hpp)
• オイラーツアー (GraphTheory/EulerTour.hpp)
• 最近共通祖先 (LCA) (GraphTheory/LowestCommonAncestor.hpp)

Code

#define PROBLEM "https://yukicoder.me/problems/no/1094"

#include "GraphTheory/EulerTour.hpp"
#include "DataStructure/BinaryIndexedTree.hpp"
#include "GraphTheory/LowestCommonAncestor.hpp"

#include <cstdio>
#include <utility>

int main() {
	int N;
	scanf("%d", &N);
	
	EulerTour<true>::Graph g(N);
	std::vector<int> A(N - 1), B(N - 1), C(N - 1);
	for (int i = 0; i < N - 1; ++i) {
		scanf("%d %d %d", &A[i], &B[i], &C[i]);
		--A[i]; --B[i];
		g[A[i]].emplace_back(B[i]);
		g[B[i]].emplace_back(A[i]);
	}
	EulerTour<true> et(g);
	
	using ll = long long;
	BinaryIndexedTree<ll> bit(et.size(), 0, [](auto x, auto y) { return x + y; });
	for (int i = 0; i < N - 1; ++i) {
		if (et.par(A[i]) != B[i]) std::swap(A[i], B[i]);
		bit.add(et.in(A[i]), C[i]);
		bit.add(et.out(A[i]), -C[i]);
	}
	LowestCommonAncestor lca(g);
	
	int Q;
	scanf("%d", &Q);
	while (Q--) {
		int s, t;
		scanf("%d %d", &s, &t);
		--s; --t;
		const int l = lca.find(s, t);
		auto func = [&](int a) {
			return bit.sum(et.in(a) + 1) - bit.sum(et.in(l) + 1);
		};
		printf("%lld\n", func(s) + func(t));
	}
}

#line 1 "Test/EulerTour.path.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/1094"

#line 1 "GraphTheory/EulerTour.hpp"



#include <vector>
#include <cassert>
#include <stack>
#include <utility>

/**
 * @brief https://tkmst201.github.io/Library/GraphTheory/EulerTour.hpp
 */
template<bool EDGE>
struct EulerTour {
	using size_type = std::size_t;
	using Graph = std::vector<std::vector<int>>;
	
private:
	int n, root;
	std::vector<int> in_, out_, par_, g_idx;
	
public:
	EulerTour(const Graph & g, int root = 0)
		: n(g.size()), root(root), in_(size(), -1), out_(size(), -1), par_(n, -1), g_idx(n << 1, -1) {
		std::stack<std::pair<int, int>> stk;
		int num = 0;
		in_[root] = num++;
		stk.emplace(root, 0);
		while (!stk.empty()) {
			const auto [u, i] = stk.top();
			stk.pop();
			if (i < g[u].size()) {
				const int v = g[u][i];
				assert(0 <= v && v < n);
				assert(v != u);
				stk.emplace(u, i + 1);
				if (v == par_[u]) g_idx[u << 1 | 1] = i;
				else {
					in_[v] = num++;
					par_[v] = u;
					g_idx[v << 1] = i;
					stk.emplace(v, 0);
				}
			}
			else {
				out_[u] = num;
				if (EDGE) ++num;
			}
		}
	}
	
	size_type size() const noexcept {
		return EDGE ? n << 1 : n;
	}
	
	int par(int k) const noexcept {
		assert(0 <= k && k < n);
		return par_[k];
	}
	
	int in(int k) const noexcept {
		assert(0 <= k && k < n);
		return in_[k];
	}
	
	int out(int k) const noexcept {
		assert(0 <= k && k < n);
		return out_[k];
	}
	
	std::pair<int, int> par_from(int k) const noexcept {
		assert(0 <= k && k < n);
		return {par_[k], g_idx[k << 1]};
	}
	
	int par_to(int k) const noexcept {
		assert(0 <= k && k < n);
		return g_idx[k << 1 | 1];
	}
};


#line 1 "DataStructure/BinaryIndexedTree.hpp"



#line 6 "DataStructure/BinaryIndexedTree.hpp"
#include <functional>

/**
 * @brief https://tkmst201.github.io/Library/DataStructure/BinaryIndexedTree.hpp
 */
template<typename T>
struct BinaryIndexedTree {
	using value_type = T;
	using const_reference = const value_type &;
	using F = std::function<value_type (const_reference, const_reference)>;
	using size_type = std::size_t;
	
private:
	size_type n;
	value_type id_elem;
	F f;
	std::vector<value_type> node;
	
public:
	BinaryIndexedTree(size_type n, const_reference id_elem, const F & f)
		: n(n), id_elem(id_elem), f(f), node(n + 1, id_elem) {}
	
	size_type size() const noexcept {
		return n;
	}
	
	void add(size_type i, const_reference x) noexcept {
		assert(i < size());
		++i;
		for (; i <= size(); i += i & -i) node[i] = f(node[i], x);
	}
	
	value_type sum(size_type i) const noexcept {
		assert(i <= size());
		value_type res = id_elem;
		for (; i > 0; i -= i & -i) res = f(node[i], res);
		return res;
	}
	
	size_type lower_bound(const_reference x) const noexcept {
		size_type res = 0;
		size_type s = id_elem, w = 1;
		while (w < size()) w <<= 1;
		for (; w > 0; w >>= 1) {
			if (res + w <= size()) {
				value_type cur = f(s, node[res + w]);
				if (cur < x) {
					res += w;
					s = cur;
				}
			}
		}
		return res;
	}
};


#line 1 "GraphTheory/LowestCommonAncestor.hpp"



#line 7 "GraphTheory/LowestCommonAncestor.hpp"
#include <cstdint>
#line 9 "GraphTheory/LowestCommonAncestor.hpp"

/**
 * @brief https://tkmst201.github.io/Library/GraphTheory/LowestCommonAncestor.hpp
 */
struct LowestCommonAncestor {
	using Graph = std::vector<std::vector<int>>;
	
private:
	int n, logn;
	std::vector<std::vector<int>> par;
	std::vector<int> depth_;
	
public:
	LowestCommonAncestor(const Graph & g, int root = 0) : n(g.size()) {
		assert(0 <= root && root < n);
		logn = 1;
		while ((1 << logn) + 1 <= n) ++logn;
		par.assign(logn, std::vector<int>(n, -1));
		depth_.assign(n, 0);
		std::stack<std::pair<int, int>> stk;
		par[0][root] = root;
		stk.emplace(root, root);
		while (!stk.empty()) {
			const auto [u, p] = stk.top();
			stk.pop();
			for (int v : g[u]) {
				if (v == p) continue;
				assert(0 <= v && v < n);
				par[0][v] = u;
				depth_[v] = depth_[u] + 1;
				stk.emplace(v, u);
			}
		}
		for (int i = 1; i < logn; ++i) {
			for (int j = 0; j < n; ++j) par[i][j] = par[i - 1][par[i - 1][j]];
		}
	}
	
	int size() const noexcept {
		return n;
	}
	
	int find(int a, int b) const noexcept {
		assert(0 <= a && a < size());
		assert(0 <= b && b < size());
		assert(par[0][a] != -1);
		assert(par[0][b] != -1);
		if (depth_[a] < depth_[b]) std::swap(a, b);
		a = parent(a, depth_[a] - depth_[b]);
		if (a == b) return a;
		for (int i = logn - 1; i >= 0; --i) {
			if (par[i][a] != par[i][b]) a = par[i][a], b = par[i][b]; 
		}
		return par[0][a];
	}
	
	int parent(int u, int k = 1) const noexcept {
		assert(0 <= u && u < size());
		assert(0 <= k && k <= depth_[u]);
		assert(par[0][u] != -1);
		for (int i = 0; i < logn; ++i) if (k >> i & 1) u = par[i][u];
		return u;
	}
	
	int depth(int u) const noexcept {
		assert(0 <= u && u < size());
		assert(par[0][u] != -1);
		return depth_[u];
	}
};


#line 6 "Test/EulerTour.path.test.cpp"

#include <cstdio>
#line 9 "Test/EulerTour.path.test.cpp"

int main() {
	int N;
	scanf("%d", &N);
	
	EulerTour<true>::Graph g(N);
	std::vector<int> A(N - 1), B(N - 1), C(N - 1);
	for (int i = 0; i < N - 1; ++i) {
		scanf("%d %d %d", &A[i], &B[i], &C[i]);
		--A[i]; --B[i];
		g[A[i]].emplace_back(B[i]);
		g[B[i]].emplace_back(A[i]);
	}
	EulerTour<true> et(g);
	
	using ll = long long;
	BinaryIndexedTree<ll> bit(et.size(), 0, [](auto x, auto y) { return x + y; });
	for (int i = 0; i < N - 1; ++i) {
		if (et.par(A[i]) != B[i]) std::swap(A[i], B[i]);
		bit.add(et.in(A[i]), C[i]);
		bit.add(et.out(A[i]), -C[i]);
	}
	LowestCommonAncestor lca(g);
	
	int Q;
	scanf("%d", &Q);
	while (Q--) {
		int s, t;
		scanf("%d %d", &s, &t);
		--s; --t;
		const int l = lca.find(s, t);
		auto func = [&](int a) {
			return bit.sum(et.in(a) + 1) - bit.sum(et.in(l) + 1);
		};
		printf("%lld\n", func(s) + func(t));
	}
}

Back to top page