Test/EulerTour.subtree.test.cpp

• View this file on GitHub
• Last update: 2021-03-12 18:03:59+09:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

• セグメント木 (DataStructure/SegmentTree.hpp)
• オイラーツアー (GraphTheory/EulerTour.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"

#include "GraphTheory/EulerTour.hpp"
#include "DataStructure/SegmentTree.hpp"

#include <vector>
#include <cstdio>

int main() {
	int N, Q;
	scanf("%d %d", &N, &Q);
	
	using ll = long long;
	std::vector<ll> A(N);
	EulerTour<false>::Graph g(N);
	for (int i = 0; i < N; ++i) scanf("%lld", &A[i]);
	std::vector<int> edge(N - 1);
	for (int i = 1; i < N; ++i) {
		scanf("%d", &edge[i - 1]);
		// par[i] = p;

		g[edge[i - 1]].emplace_back(i);
	}
	
	EulerTour<false> et(g, 0);
	
	std::vector<ll> tv(et.size());
	for (int i = 0; i < N; ++i) tv[et.in(i)] = A[i];
	SegmentTree<ll> seg(tv, 0, [](ll x, ll y) { return x + y; });
	
	while (Q--) {
		int q;
		scanf("%d", &q);
		if (q == 0) {
			int u, x;
			scanf("%d %d", &u, &x);
			const int idx = et.in(u);
			seg.set(idx, seg.get(idx) + x);
		}
		else {
			int u;
			scanf("%d", &u);
			printf("%lld\n", seg.fold(et.in(u), et.out(u)));
		}
	}
	
	// par* test

	for (int i = 0; i < N; ++i) {
		if (i == 0) {
			assert(et.par(i) == -1);
			assert(et.par_from(i) == std::make_pair(-1, -1));
			assert(et.par_to(i) == -1);
		}
		else {
			const int p = et.par(i);
			auto [p1, pfidx] = et.par_from(i);
			assert(p == p1);
			assert(pfidx < g[p].size());
			assert(g[p][pfidx] == i);
			assert(et.par_to(i) == -1);
		}
	}
	
	for (int i = 1; i < N; ++i) g[i].emplace_back(edge[i - 1]);
	et = EulerTour<false>(g);
	for (int i = 0; i < N; ++i) {
		if (i == 0) {
			assert(et.par(i) == -1);
			assert(et.par_from(i) == std::make_pair(-1, -1));
			assert(et.par_to(i) == -1);
		}
		else {
			const int p = et.par(i);
			auto [p1, pfidx] = et.par_from(i);
			assert(p == p1);
			assert(pfidx < g[p].size());
			assert(g[p][pfidx] == i);
			auto ptidx = et.par_to(i);
			assert(ptidx < g[i].size());
			assert(g[i][ptidx] == p);
		}
	}
}

#line 1 "Test/EulerTour.subtree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"

#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/SegmentTree.hpp"



#line 5 "DataStructure/SegmentTree.hpp"
#include <algorithm>
#line 7 "DataStructure/SegmentTree.hpp"
#include <functional>

/**
 * @brief https://tkmst201.github.io/Library/DataStructure/SegmentTree.hpp
 */
template<typename T>
struct SegmentTree {
	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, n_;
	value_type id_elem;
	F f;
	std::vector<value_type> node;
	
public:
	SegmentTree() = default;
	
	SegmentTree(size_type n, const_reference id_elem, const F & f)
		: n(n), id_elem(id_elem), f(f) {
		n_ = 1;
		while (n_ < n) n_ <<= 1;
		node.assign(2 * n_, id_elem);
	}
	
	SegmentTree(const std::vector<value_type> & v, const_reference id_elem, const F & f)
		: SegmentTree(v.size(), id_elem, f) {
		for (size_type i = 0; i < v.size(); ++i) node[i + n_] = v[i];
		for (size_type i = n_ - 1; i > 0; --i) node[i] = f(node[i << 1], node[i << 1 | 1]);
	}
	
	size_type size() const noexcept {
		return n;
	}
	
	void set(size_type i, const_reference x) noexcept {
		assert(i < size());
		node[i += n_] = x;
		while (i > 1) {
			i >>= 1;
			node[i] = f(node[i << 1], node[i << 1 | 1]);
		}
	}
	
	const_reference get(size_type i) const noexcept {
		assert(i < size());
		return node[i + n_];
	}
	
	value_type fold(size_type l, size_type r) const noexcept {
		assert(l <= r);
		assert(r <= size());
		value_type lv = id_elem, rv = id_elem;
		for (l += n_, r += n_; l < r; l >>= 1, r >>= 1) {
			if (l & 1) lv = f(lv, node[l++]);
			if (r & 1) rv = f(node[r - 1], rv);
		}
		return f(lv, rv);
	}
	
	value_type fold_all() const noexcept {
		return node[1];
	}
	
	size_type max_right(size_type l, std::function<bool (const_reference)> g) const noexcept {
		assert(l <= size());
		assert(g(id_elem));
		if (l == size()) return size();
		l += n_;
		value_type sum = id_elem;
		while (true) {
			while (~l & 1) l >>= 1;
			const value_type nex_sum = f(sum, node[l]);
			if (g(nex_sum)) { sum = nex_sum; ++l; }
			else break;
			if ((l & -l) == l) return size();
		}
		while (l < n_) {
			const value_type nex_sum = f(sum, node[l << 1]);
			l <<= 1;
			if (g(nex_sum)) { sum = nex_sum; l |= 1; }
		}
		return l - n_;
	}
	
	size_type min_left(size_type r, std::function<bool (const_reference)> g) const noexcept {
		assert(r <= size());
		assert(g(id_elem));
		if (r == 0) return 0;
		r += n_;
		value_type sum = id_elem;
		while (true) {
			--r;
			while (r > 1 && (r & 1)) r >>= 1;
			const value_type nex_sum = f(node[r], sum);
			if (g(nex_sum)) sum = nex_sum;
			else break;
			if ((r & -r) == r) return 0;
		}
		while (r < n_) {
			const value_type nex_sum = f(node[r << 1 | 1], sum);
			r <<= 1;
			if (!g(nex_sum)) r |= 1;
			else sum = nex_sum;
		}
		return r + 1 - n_;
	}
};


#line 5 "Test/EulerTour.subtree.test.cpp"

#line 7 "Test/EulerTour.subtree.test.cpp"
#include <cstdio>

int main() {
	int N, Q;
	scanf("%d %d", &N, &Q);
	
	using ll = long long;
	std::vector<ll> A(N);
	EulerTour<false>::Graph g(N);
	for (int i = 0; i < N; ++i) scanf("%lld", &A[i]);
	std::vector<int> edge(N - 1);
	for (int i = 1; i < N; ++i) {
		scanf("%d", &edge[i - 1]);
		// par[i] = p;

		g[edge[i - 1]].emplace_back(i);
	}
	
	EulerTour<false> et(g, 0);
	
	std::vector<ll> tv(et.size());
	for (int i = 0; i < N; ++i) tv[et.in(i)] = A[i];
	SegmentTree<ll> seg(tv, 0, [](ll x, ll y) { return x + y; });
	
	while (Q--) {
		int q;
		scanf("%d", &q);
		if (q == 0) {
			int u, x;
			scanf("%d %d", &u, &x);
			const int idx = et.in(u);
			seg.set(idx, seg.get(idx) + x);
		}
		else {
			int u;
			scanf("%d", &u);
			printf("%lld\n", seg.fold(et.in(u), et.out(u)));
		}
	}
	
	// par* test

	for (int i = 0; i < N; ++i) {
		if (i == 0) {
			assert(et.par(i) == -1);
			assert(et.par_from(i) == std::make_pair(-1, -1));
			assert(et.par_to(i) == -1);
		}
		else {
			const int p = et.par(i);
			auto [p1, pfidx] = et.par_from(i);
			assert(p == p1);
			assert(pfidx < g[p].size());
			assert(g[p][pfidx] == i);
			assert(et.par_to(i) == -1);
		}
	}
	
	for (int i = 1; i < N; ++i) g[i].emplace_back(edge[i - 1]);
	et = EulerTour<false>(g);
	for (int i = 0; i < N; ++i) {
		if (i == 0) {
			assert(et.par(i) == -1);
			assert(et.par_from(i) == std::make_pair(-1, -1));
			assert(et.par_to(i) == -1);
		}
		else {
			const int p = et.par(i);
			auto [p1, pfidx] = et.par_from(i);
			assert(p == p1);
			assert(pfidx < g[p].size());
			assert(g[p][pfidx] == i);
			auto ptidx = et.par_to(i);
			assert(ptidx < g[i].size());
			assert(g[i][ptidx] == p);
		}
	}
}

Back to top page