verify/yosupo/yosupo-vertex_add_path_sum-new-hld.test.cpp

• View this file on GitHub
• Last update: 2022-07-21 16:12:33-04:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_path_sum

Depends on

• library/contest/template-minimal.hpp
• library/data-structures/1d-range-queries/general-full-segment-tree.hpp
• library/graphs/heavy-light-decomposition.hpp

Code

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

#include "../../library/contest/template-minimal.hpp"
#include "../../library/graphs/heavy-light-decomposition.hpp"
#include "../../library/data-structures/1d-range-queries/general-full-segment-tree.hpp"

int main() {
	using namespace std;
	cin.tie(0)->sync_with_stdio(false);
	int n, q;
	struct Node {
		long long v;
		int size;
		
		Node() : v(0), size(1) {}
		Node(long long _v, int _size) : v(_v), size(_size) {}
	};
	cin >> n >> q;
	vector<int> a(n);
	for (int i = 0; i < n; i++)
		cin >> a[i];
	vector<vector<int>> graph(n);
	for (int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		graph[u].push_back(v);
		graph[v].push_back(u);
	}
	HeavyLightDecomposition H(graph);
	auto seg = get_lazy_segment_tree(
		vector<Node>(n),
		Node(), 0LL,
		[&](Node x, Node y) { return Node{x.v + y.v, x.size + y.size}; },
		[&](Node x, long long y) { return Node{x.v + y * x.size, x.size}; },
		[&](long long x, long long y) { return x + y; }
	);
	for (int i = 0; i < n; i++)
		H.path_query(i, i, true, [&](int l, int r) {
			seg.add(l, r, a[i]);
		});
	while (q--) {
		int t;
		cin >> t;
		if (t == 0) {
			int p, x;
			cin >> p >> x;
			H.path_query(p, p, true, [&](int l, int r) {
				seg.add(l, r, x);
			});
		} else	{
			int u, v;
			cin >> u >> v;
			long long ans = 0;
			H.path_query(u, v, true, [&](int l, int r) {
				ans += seg.sum(l, r).v;
			});
			cout << ans << '\n';
		}
	}
	return 0;
}

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


#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <iostream>
#include <iomanip>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>

using namespace std;

template <class G> 
struct HeavyLightDecomposition {
private:
	void dfs_sz(int cur) {
		size[cur] = 1;
		for (auto& dst : g[cur]) {
			if (dst == par[cur]) {
				if (g[cur].size() >= 2 && int(dst) == int(g[cur][0]))
					std::swap(g[cur][0], g[cur][1]);
				else
					continue;
			}
			depth[dst] = depth[cur] + 1;
			par[dst] = cur;
			dfs_sz(dst);
			size[cur] += size[dst];
			if (size[dst] > size[g[cur][0]]) {
				std::swap(dst, g[cur][0]);
			}
		}
	}

	void dfs_hld(int cur) {
		down[cur] = id++;
		for (auto dst : g[cur]) {
			if (dst == par[cur]) continue;
			nxt[dst] = (int(dst) == int(g[cur][0]) ? nxt[cur] : int(dst));
			dfs_hld(dst);
		}
		up[cur] = id;
	}

	// [u, v)

	std::vector<std::pair<int, int>> ascend(int u, int v) const {
		std::vector<std::pair<int, int>> res;
		while (nxt[u] != nxt[v]) {
			res.emplace_back(down[u], down[nxt[u]]);
			u = par[nxt[u]];
		}
		if (u != v) res.emplace_back(down[u], down[v] + 1);
		return res;
	}

	// (u, v]

	std::vector<std::pair<int, int>> descend(int u, int v) const {
		if (u == v) return {};
		if (nxt[u] == nxt[v]) return {{down[u] + 1, down[v]}};
		auto res = descend(u, par[nxt[v]]);
		res.emplace_back(down[nxt[v]], down[v]);
		return res;
	}

public:
	G g;
	int id;
	std::vector<int> size, depth, down, up, nxt, par;

	HeavyLightDecomposition(G& _g, std::vector<int> roots = {0})
			: g(_g),
				id(0),
				size(g.size(), 0),
				depth(g.size(), 0),
				down(g.size(), -1),
				up(g.size(), -1),
				nxt(g.size(), 0), 
				par(g.size(), 0) {
		for (int root : roots) {
			par[root] = nxt[root] = root;
			dfs_sz(root);
			dfs_hld(root);
		}
	}

	void build(std::vector<int> roots) {
		for (int root : roots) {
			par[root] = nxt[root] = root;
			dfs_sz(root);
			dfs_hld(root);
		}
	}

	// [l, r], inclusive for subtree

	std::pair<int, int> idx(int i) const { return {down[i], up[i] - 1}; }

	template <class F>
	void path_query(int u, int v, bool vertex, const F& f) {
		int l = lca(u, v);
		for (auto&& [a, b] : ascend(u, l)) {
			int s = a + 1, t = b;
			s > t ? f(t, s - 1) : f(s, t - 1);
		}
		if (vertex) f(down[l], down[l]);
		for (auto&& [a, b] : descend(l, v)) {
			int s = a, t = b + 1;
			s > t ? f(t, s - 1) : f(s, t - 1);
		}
	}

	template <class F>
	void path_noncommutative_query(int u, int v, bool vertex, const F& f) {
		int l = lca(u, v);
		for (auto&& [a, b] : ascend(u, l)) f(a + 1, b - 1);
		if (vertex) f(down[l], down[l]);
		for (auto&& [a, b] : descend(l, v)) f(a, b);
	}

	template <class F>
	void subtree_query(int u, bool vertex, const F& f) {
		f(down[u] + int(!vertex), up[u] - 1);
	}

	int lca(int a, int b) {
		while (nxt[a] != nxt[b]) {
			if (down[a] < down[b]) std::swap(a, b);
			a = par[nxt[a]];
		}
		return depth[a] < depth[b] ? a : b;
	}

	int dist(int a, int b) { return depth[a] + depth[b] - depth[lca(a, b)] * 2; }
};

template <class D, class L, class OpDD, class OpDL, class OpLL> struct LazySegmentTree {
	D e_d;
	L e_l;
	OpDD op_dd; 
	OpDL op_dl;
	OpLL op_ll;
	int sz, lg;  
	std::vector<D> d;
	std::vector<L> lz;

	void init(const std::vector<D>& v) {
		int n = int(v.size());
		lg = 1;
		while ((1 << lg) < n) lg++;
		sz = 1 << lg;
		d = std::vector<D>(2 * sz, e_d);
		lz = std::vector<L>(2 * sz, e_l);
		for (int i = 0; i < n; i++) d[sz + i] = v[i];
		for (int i = sz - 1; i >= 0; i--) {
			update(i);
		}
	}

	LazySegmentTree(const std::vector<D>& v,
			D _e_d,
			L _e_l,
			OpDD _op_dd,
			OpDL _op_dl,
			OpLL _op_ll)
		: e_d(_e_d), e_l(_e_l), op_dd(_op_dd), op_dl(_op_dl), op_ll(_op_ll) {
		init(v);
	}

	void all_add(int k, L x) {
		d[k] = op_dl(d[k], x);
		if (k < sz) lz[k] = op_ll(lz[k], x);
	}

	void push(int k) {
		all_add(2 * k, lz[k]);
		all_add(2 * k + 1, lz[k]);
		lz[k] = e_l;
	}

	void update(int k) { d[k] = op_dd(d[2 * k], d[2 * k + 1]); }

	void set(int p, D x) {
		p += sz;
		for (int i = lg; i >= 1; i--) push(p >> i);
		d[p] = x;
		for (int i = 1; i <= lg; i++) update(p >> i);
	}

	void add(int a, int b, L x, int l, int r, int k) {
		if (b <= l || r <= a) return;
		if (a <= l && r <= b) {
			all_add(k, x);
			return;
		}
		push(k);
		int mid = (l + r) / 2;
		add(a, b, x, l, mid, 2 * k);
		add(a, b, x, mid, r, 2 * k + 1);
		update(k);
	}

	void add(int a, int b, L x) {
		++b;
		a += sz;
		b += sz;
		for (int i = lg; i >= 1; i--) {
			if (((a >> i) << i) != a) push(a >> i);
			if (((b >> i) << i) != b) push((b - 1) >> i);
		}
		{
			int a2 = a, b2 = b;
			while (a < b) {
				if (a & 1) all_add(a++, x);
				if (b & 1) all_add(--b, x);
				a >>= 1;
				b >>= 1;
			}
			a = a2;
			b = b2;
		}
		for (int i = 1; i <= lg; i++) {
			if (((a >> i) << i) != a) update(a >> i);
			if (((b >> i) << i) != b) update((b - 1) >> i);
		}
	}

	D single(int p) {
		p += sz;
		for (int i = lg; i >= 1; i--) push(p >> i);
		return d[p];
	}

	D sum(int a, int b, int l, int r, int k) {
		if (b <= l || r <= a) return e_d;
		if (a <= l && r <= b) return d[k];
		push(k);
		int mid = (l + r) / 2;
		return op_dd(sum(a, b, l, mid, 2 * k), sum(a, b, mid, r, 2 * k + 1));
	}

	D sum(int a, int b) {
		++b;
		if (a == b) return e_d;
		a += sz;
		b += sz;
		for (int i = lg; i >= 1; i--) {
			if (((a >> i) << i) != a) push(a >> i);
			if (((b >> i) << i) != b) push((b - 1) >> i);
		}
		D sml = e_d, smr = e_d;
		while (a < b) {
			if (a & 1) sml = op_dd(sml, d[a++]);
			if (b & 1) smr = op_dd(d[--b], smr);
			a >>= 1;
			b >>= 1;
		}
		return op_dd(sml, smr);
	}

	D all_sum() const { return d[1]; }
};

template <class D, class L, class OpDD, class OpDL, class OpLL>
LazySegmentTree<D, L, OpDD, OpDL, OpLL> get_lazy_segment_tree(std::vector<D> v,
										D e_d,
										L e_l,
										OpDD op_dd,
										OpDL op_dl,
										OpLL op_ll) {
	return LazySegmentTree<D, L, OpDD, OpDL, OpLL>(v, e_d, e_l, op_dd, op_dl, op_ll);
}

int main() {
	using namespace std;
	cin.tie(0)->sync_with_stdio(false);
	int n, q;
	struct Node {
		long long v;
		int size;
		
		Node() : v(0), size(1) {}
		Node(long long _v, int _size) : v(_v), size(_size) {}
	};
	cin >> n >> q;
	vector<int> a(n);
	for (int i = 0; i < n; i++)
		cin >> a[i];
	vector<vector<int>> graph(n);
	for (int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		graph[u].push_back(v);
		graph[v].push_back(u);
	}
	HeavyLightDecomposition H(graph);
	auto seg = get_lazy_segment_tree(
		vector<Node>(n),
		Node(), 0LL,
		[&](Node x, Node y) { return Node{x.v + y.v, x.size + y.size}; },
		[&](Node x, long long y) { return Node{x.v + y * x.size, x.size}; },
		[&](long long x, long long y) { return x + y; }
	);
	for (int i = 0; i < n; i++)
		H.path_query(i, i, true, [&](int l, int r) {
			seg.add(l, r, a[i]);
		});
	while (q--) {
		int t;
		cin >> t;
		if (t == 0) {
			int p, x;
			cin >> p >> x;
			H.path_query(p, p, true, [&](int l, int r) {
				seg.add(l, r, x);
			});
		} else	{
			int u, v;
			cin >> u >> v;
			long long ans = 0;
			H.path_query(u, v, true, [&](int l, int r) {
				ans += seg.sum(l, r).v;
			});
			cout << ans << '\n';
		}
	}
	return 0;
}

Back to top page