test/HLD_subtree_edge.test.cpp

• View this file on GitHub
• Last update: 2021-03-21 11:17:59+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/Rupc/2667?year=2015

Depends on

• DataStructure/LazySegmentTree.cpp
• Graph/HeavyLightDecomposition.cpp
• atcoder/internal_bit.hpp
• atcoder/lazysegtree.hpp

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/Rupc/2667?year=2015"
#include "./../Graph/HeavyLightDecomposition.cpp"
#include "./../DataStructure/LazySegmentTree.cpp"
#include <iostream>
using namespace std;
using ll = long long;

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

	int n, q;
	cin >> n >> q;
	HLD g(n);
	for (int i = 0; i < n - 1; ++i) {
		int a, b;
		cin >> a >> b;
		g.add_edge(a, b);
	}
	g.build(0);

	RangeAddRangeSum<ll, ll> seg(vector<S_sum<ll>>(n, 0));
	while (q--) {
		int com;
		cin >> com;
		if (com == 0) {
			int u, v;
			cin >> u >> v;
			ll ans = 0;
			g.each_edge(u, v, [&](int l, int r) { ans += seg.prod(l, r).value; });
			cout << ans << '\n';
		} else {
			int v;
			ll x;
			cin >> v >> x;
			g.each_subtree_edge(v, [&](int l, int r) { seg.apply(l, r, x); });
		}
	}
}

#line 1 "test/HLD_subtree_edge.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/Rupc/2667?year=2015"
#line 2 "Graph/HeavyLightDecomposition.cpp"
#include <vector>
#include <cassert>

class HLD {
	int n;
	std::vector<std::vector<int>> graph;
	std::vector<int> parent, size;
	int k;
	std::vector<int> head, hld, index, out_index;
	bool builded = false;

	int calc_size(int v, int p, int d) {
		parent[v] = p;
		size[v] = 1;
		for (int u : graph[v]) {
			if (u != p) {
				size[v] += calc_size(u, v, d + 1);
			}
		}
		return size[v];
	}
	void rec(int v, int p, int root) {
		head[v] = root;
		index[v] = k;
		hld[k++] = v;

		int heavy_vertex = -1, max_size = 0;
		for (int u : graph[v]) {
			if (u != p && max_size < size[u]) {
				max_size = size[u];
				heavy_vertex = u;
			}
		}
		if (heavy_vertex != -1) {
			rec(heavy_vertex, v, root);
			for (int u : graph[v]) {
				if (u != heavy_vertex && u != p) {
					rec(u, v, u);
				}
			}
		}
		out_index[v] = k;
	}

public:
	HLD(int _n) : n(_n), graph(_n) {}
	HLD(const std::vector<std::vector<int>>& _graph) : n(_graph.size()), graph(_graph) {}
	void add_edge(int u, int v) {
		graph[u].push_back(v);
		graph[v].push_back(u);
		builded = false;
	}
	void build(int root) {
		parent.assign(n, -1);
		size.assign(n, 0);
		calc_size(root, -1, 1);
		k = 0;
		head.assign(n, 0);
		hld.assign(n, 0);
		index.assign(n, 0);
		out_index.assign(n, 0);
		rec(root, -1, root);
		builded = true;
	}
	const std::vector<std::vector<int>>& get_graph() const {
		assert(builded);
		return graph;
	}
	const std::vector<int>& get_parent() const {
		assert(builded);
		return parent;
	}
	const std::vector<int>& get_size() const {
		assert(builded);
		return size;
	}
	const std::vector<int>& get_head() const {
		assert(builded);
		return head;
	}
	const std::vector<int>& get_hld() const {
		assert(builded);
		return hld;
	}
	const std::vector<int>& get_index() const {
		assert(builded);
		return index;
	}
	const std::vector<int>& get_out_index() const {
		assert(builded);
		return out_index;
	}
	int operator[](int v) const {
		assert(builded);
		return index[v];
	}

	template <class F> void each_vertex(int v, int u, F f) const {
		assert(builded);
		while (true) {
			if (index[v] > index[u]) std::swap(v, u);
			if (head[v] != head[u]) {
				f(index[head[u]], index[u] + 1);
				u = parent[head[u]];
			} else {
				f(index[v], index[u] + 1);
				break;
			}
		}
	}
	template <class F> void each_subtree_vertex(int v, F f) const {
		assert(builded);
		f(index[v], out_index[v]);
	}
	template <class F> void each_edge(int v, int u, F f) const {
		assert(builded);
		while (true) {
			if (index[v] > index[u]) std::swap(v, u);
			if (head[v] != head[u]) {
				f(index[head[u]], index[u] + 1);
				u = parent[head[u]];
			} else {
				if (v != u) f(index[v] + 1, index[u] + 1);
				break;
			}
		}
	}
	template <class F> void each_subtree_edge(int v, F f) const {
		assert(builded);
		f(index[v] + 1, out_index[v]);
	}
	std::vector<std::pair<int, int>> query_vertex(int u, int v) const {
		assert(builded);
		std::vector<std::pair<int, int>> result;
		each_vertex(u, v, [&](int a, int b) { result.emplace_back(a, b); });
		return result;
	}
	std::pair<int, int> query_subtree_vertex(int v) const {
		assert(builded);
		std::pair<int, int> result;
		each_subtree_vertex(v, [&](int a, int b) { result = {a, b}; });
		return result;
	}
	std::vector<std::pair<int, int>> query_edge(int u, int v) const {
		assert(builded);
		std::vector<std::pair<int, int>> result;
		each_edge(u, v, [&](int a, int b) { result.emplace_back(a, b); });
		return result;
	}
	std::pair<int, int> query_subtree_edge(int v) const {
		assert(builded);
		std::pair<int, int> result;
		each_subtree_edge(v, [&](int a, int b) { result = {a, b}; });
		return result;
	}
	int lca(int u, int v) const {
		while (true) {
			if (index[u] > index[v]) std::swap(u, v);
			if (head[u] != head[v]) {
				v = parent[head[v]];
			} else {
				return u;
			}
		}
	}
};
#line 1 "atcoder/lazysegtree.hpp"



#line 1 "atcoder/internal_bit.hpp"



#ifdef _MSC_VER
#include <intrin.h>
#endif

namespace atcoder {

	namespace internal {

		// @param n `0 <= n`
		// @return minimum non-negative `x` s.t. `n <= 2**x`
		int ceil_pow2(int n) {
			int x = 0;
			while ((1U << x) < (unsigned int)(n))
				x++;
			return x;
		}

		// @param n `1 <= n`
		// @return minimum non-negative `x` s.t. `(n & (1 << x)) != 0`
		int bsf(unsigned int n) {
#ifdef _MSC_VER
			unsigned long index;
			_BitScanForward(&index, n);
			return index;
#else
			return __builtin_ctz(n);
#endif
		}

	}  // namespace internal

}  // namespace atcoder


#line 5 "atcoder/lazysegtree.hpp"
#include <algorithm>
#line 7 "atcoder/lazysegtree.hpp"
#include <iostream>
#line 9 "atcoder/lazysegtree.hpp"
namespace atcoder {

	template <class S, S (*op)(S, S), S (*e)(), class F, S (*mapping)(F, S),
	          F (*composition)(F, F), F (*id)()>
	struct lazy_segtree {
	public:
		lazy_segtree() : lazy_segtree(0) {}
		lazy_segtree(int n) : lazy_segtree(std::vector<S>(n, e())) {}
		lazy_segtree(const std::vector<S>& v) : _n(int(v.size())) {
			log = internal::ceil_pow2(_n);
			size = 1 << log;
			d = std::vector<S>(2 * size, e());
			lz = std::vector<F>(size, id());
			for (int i = 0; i < _n; i++) d[size + i] = v[i];
			for (int i = size - 1; i >= 1; i--) {
				update(i);
			}
		}

		void set(int p, S x) {
			assert(0 <= p && p < _n);
			p += size;
			for (int i = log; i >= 1; i--) push(p >> i);
			d[p] = x;
			for (int i = 1; i <= log; i++) update(p >> i);
		}

		S get(int p) {
			assert(0 <= p && p < _n);
			p += size;
			for (int i = log; i >= 1; i--) push(p >> i);
			return d[p];
		}

		S operator[](int p) {
			return get(p);
		}

		S prod(int l, int r) {
			assert(0 <= l && l <= r && r <= _n);
			if (l == r) return e();

			l += size;
			r += size;

			for (int i = log; i >= 1; i--) {
				if (((l >> i) << i) != l) push(l >> i);
				if (((r >> i) << i) != r) push(r >> i);
			}

			S sml = e(), smr = e();
			while (l < r) {
				if (l & 1) sml = op(sml, d[l++]);
				if (r & 1) smr = op(d[--r], smr);
				l >>= 1;
				r >>= 1;
			}

			return op(sml, smr);
		}

		S operator()(int l, int r) {
			return prod(l, r);
		}

		S all_prod() {
			return d[1];
		}

		void apply(int p, F f) {
			assert(0 <= p && p < _n);
			p += size;
			for (int i = log; i >= 1; i--) push(p >> i);
			d[p] = mapping(f, d[p]);
			for (int i = 1; i <= log; i++) update(p >> i);
		}
		void apply(int l, int r, F f) {
			assert(0 <= l && l <= r && r <= _n);
			if (l == r) return;

			l += size;
			r += size;

			for (int i = log; i >= 1; i--) {
				if (((l >> i) << i) != l) push(l >> i);
				if (((r >> i) << i) != r) push((r - 1) >> i);
			}

			{
				int l2 = l, r2 = r;
				while (l < r) {
					if (l & 1) all_apply(l++, f);
					if (r & 1) all_apply(--r, f);
					l >>= 1;
					r >>= 1;
				}
				l = l2;
				r = r2;
			}

			for (int i = 1; i <= log; i++) {
				if (((l >> i) << i) != l) update(l >> i);
				if (((r >> i) << i) != r) update((r - 1) >> i);
			}
		}

		template <bool (*g)(S)> int max_right(int l) {
			return max_right(l, [](S x) { return g(x); });
		}
		template <class G> int max_right(int l, G g) {
			assert(0 <= l && l <= _n);
			assert(g(e()));
			if (l == _n) return _n;
			l += size;
			for (int i = log; i >= 1; i--) push(l >> i);
			S sm = e();
			do {
				while (l % 2 == 0) l >>= 1;
				if (!g(op(sm, d[l]))) {
					while (l < size) {
						push(l);
						l = (2 * l);
						if (g(op(sm, d[l]))) {
							sm = op(sm, d[l]);
							l++;
						}
					}
					return l - size;
				}
				sm = op(sm, d[l]);
				l++;
			} while ((l & -l) != l);
			return _n;
		}

		template <bool (*g)(S)> int min_left(int r) {
			return min_left(r, [](S x) { return g(x); });
		}
		template <class G> int min_left(int r, G g) {
			assert(0 <= r && r <= _n);
			assert(g(e()));
			if (r == 0) return 0;
			r += size;
			for (int i = log; i >= 1; i--) push((r - 1) >> i);
			S sm = e();
			do {
				r--;
				while (r > 1 && (r % 2)) r >>= 1;
				if (!g(op(d[r], sm))) {
					while (r < size) {
						push(r);
						r = (2 * r + 1);
						if (g(op(d[r], sm))) {
							sm = op(d[r], sm);
							r--;
						}
					}
					return r + 1 - size;
				}
				sm = op(d[r], sm);
			} while ((r & -r) != r);
			return 0;
		}

		std::vector<S> to_a() {
			std::vector<S> res(_n);
			for (int i = 0; i < _n; ++i) {
				res[i] = get(i);
			}
			return res;
		}

	private:
		int _n, size, log;
		std::vector<S> d;
		std::vector<F> lz;

		void update(int k) {
			d[k] = op(d[2 * k], d[2 * k + 1]);
		}
		void all_apply(int k, F f) {
			d[k] = mapping(f, d[k]);
			if (k < size) lz[k] = composition(f, lz[k]);
		}
		void push(int k) {
			all_apply(2 * k, lz[k]);
			all_apply(2 * k + 1, lz[k]);
			lz[k] = id();
		}
	};

}  // namespace atcoder


#line 4 "DataStructure/LazySegmentTree.cpp"
#include <limits>

namespace internal {
	template <class T> struct S_sum {
		T value, size;
		S_sum(T v, T s = 1) : value(v), size(s) {}
	};

	template <class S> constexpr S constant_min() {
		return std::numeric_limits<S>::min();
	}
	template <class S> constexpr S constant_max() {
		return std::numeric_limits<S>::max();
	}
	template <class S> constexpr S constant_zero() {
		return static_cast<S>(0);
	}
	template <class T> constexpr S_sum<T> constant_zero_sum() {
		return {0, 0};
	}

	template <class S> constexpr S op_max(S x, S y) {
		return std::max(x, y);
	}
	template <class S> constexpr S op_min(S x, S y) {
		return std::min(x, y);
	}
	template <class T> constexpr S_sum<T> op_sum(S_sum<T> x, S_sum<T> y) {
		return {x.value + y.value, x.size + y.size};
	}

	template <class S, class F> constexpr S mapping_add(F f, S x) {
		return f + x;
	}
	template <class T, class F> constexpr S_sum<T> mapping_add_sum(F f, S_sum<T> x) {
		return {x.value + f * x.size, x.size};
	}
	template <class S, class F> constexpr S mapping_update(F f, S x) {
		return f == constant_max<F>() ? x : f;
	}
	template <class T, class F> constexpr S_sum<T> mapping_update_sum(F f, S_sum<T> x) {
		if (f != constant_max<F>()) x.value = f * x.size;
		return x;
	}

	template <class F> constexpr F composition_add(F f, F g) {
		return f + g;
	}
	template <class F> constexpr F composition_update(F f, F g) {
		return f == constant_max<F>() ? g : f;
	}
}  // namespace internal

using internal::S_sum;

template <class S, class F>
using RangeAddRangeMax =
    atcoder::lazy_segtree<S, internal::op_max<S>, internal::constant_min<S>, F,
                          internal::mapping_add<S, F>, internal::composition_add<F>,
                          internal::constant_zero<F>>;
template <class S, class F>
using RangeAddRangeMin =
    atcoder::lazy_segtree<S, internal::op_min<S>, internal::constant_max<S>, F,
                          internal::mapping_add<S, F>, internal::composition_add<F>,
                          internal::constant_zero<F>>;
template <class T, class F>
using RangeAddRangeSum =
    atcoder::lazy_segtree<internal::S_sum<T>, internal::op_sum<T>,
                          internal::constant_zero_sum<T>, F, internal::mapping_add_sum<T, F>,
                          internal::composition_add<F>, internal::constant_zero<F>>;
template <class S, class F>
using RangeUpdateRangeMax =
    atcoder::lazy_segtree<S, internal::op_max<S>, internal::constant_min<S>, F,
                          internal::mapping_update<S, F>, internal::composition_update<F>,
                          internal::constant_max<F>>;
template <class S, class F>
using RangeUpdateRangeMin =
    atcoder::lazy_segtree<S, internal::op_min<S>, internal::constant_max<S>, F,
                          internal::mapping_update<S, F>, internal::composition_update<F>,
                          internal::constant_max<F>>;
template <class T, class F>
using RangeUpdateRangeSum =
    atcoder::lazy_segtree<internal::S_sum<T>, internal::op_sum<T>,
                          internal::constant_zero_sum<T>, F,
                          internal::mapping_update_sum<T, F>, internal::composition_update<F>,
                          internal::constant_max<F>>;
#line 5 "test/HLD_subtree_edge.test.cpp"
using namespace std;
using ll = long long;

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

	int n, q;
	cin >> n >> q;
	HLD g(n);
	for (int i = 0; i < n - 1; ++i) {
		int a, b;
		cin >> a >> b;
		g.add_edge(a, b);
	}
	g.build(0);

	RangeAddRangeSum<ll, ll> seg(vector<S_sum<ll>>(n, 0));
	while (q--) {
		int com;
		cin >> com;
		if (com == 0) {
			int u, v;
			cin >> u >> v;
			ll ans = 0;
			g.each_edge(u, v, [&](int l, int r) { ans += seg.prod(l, r).value; });
			cout << ans << '\n';
		} else {
			int v;
			ll x;
			cin >> v >> x;
			g.each_subtree_edge(v, [&](int l, int r) { seg.apply(l, r, x); });
		}
	}
}

Back to top page