test/tree/hld/yosupo-Vertex-Add-Path-Sum.test.cpp

• View this file on GitHub
• Last update: 2023-08-13 14:16:40+08:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_path_sum

Depends on

• library/data-structure/fenwick.hpp
• library/data-structure/sparse-table.hpp
• library/tree/hld.hpp

Code

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

#include <iostream>

#include "../../../library/data-structure/fenwick.hpp"

#include "../../../library/tree/hld.hpp"

using namespace std;
using namespace felix;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	int n, q;
	cin >> n >> q;
	vector<int> a(n);
	for(int i = 0; i < n; i++) {
		cin >> a[i];
	}
	HLD hld(n);
	for(int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		hld.add_edge(u, v);
	}
	hld.build(0);
	fenwick<long long> fenw(n);
	for(int i = 0; i < n; i++) {
		fenw.add(hld.id[i], a[i]);
		fenw.add(hld.id[i] + hld.subtree_size[i], -a[i]);
	}
	while(q--) {
		int type, u, v;
		cin >> type >> u >> v;
		if(type == 0) {
			fenw.add(hld.id[u], v);
			fenw.add(hld.id[u] + hld.subtree_size[u], -v);
		} else {
			int z = hld.get_lca(u, v);
			long long ans = fenw.get(hld.id[u] + 1) + fenw.get(hld.id[v] + 1) - fenw.get(hld.id[z] + 1);
			if(hld.parent[z] != -1) {
				ans -= fenw.get(hld.id[hld.parent[z]] + 1);
			}
			cout << ans << "\n";
		}
	}
	return 0;
}

#line 1 "test/tree/hld/yosupo-Vertex-Add-Path-Sum.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_path_sum"

#include <iostream>

#line 2 "library/data-structure/fenwick.hpp"
#include <vector>
#include <cassert>

namespace felix {

template<class S>
struct fenwick {
public:
	fenwick() : n(0) {}
	explicit fenwick(int _n) : n(_n), data(_n) {}

	void add(int p, S x) {
		for(int i = p + 1; i <= n; i += i & -i) {
			data[i - 1] += x;
		}
	}

	// [0, p)
	S get(int p) const {
		auto ans = S();
		for(int i = p; i > 0; i -= i & -i) {
			ans += data[i - 1];
		}
		return ans;
	}

	// [l, r)
	S sum(int l, int r) const { return get(r) - get(l); }

	// 0-based
	int kth(S k) const {
		int x = 0;
		for(int i = 1 << std::__lg(n); i > 0; i >>= 1) {
			if (x + i <= n && k >= data[x + i - 1]) {
				x += i;
				k -= data[x - 1];
			}
		}
		return x;
	}

private:
	int n;
	std::vector<S> data;
};

} // namespace felix
#line 3 "library/tree/hld.hpp"
#include <array>

#line 5 "library/tree/hld.hpp"
#include <algorithm>

#include <cmath>

#line 4 "library/data-structure/sparse-table.hpp"

namespace felix {

template<class S, S (*op)(S, S)>
struct sparse_table {
public:
	sparse_table() {}
	explicit sparse_table(const std::vector<S>& a) {
		n = (int) a.size();
		int max_log = std::__lg(n) + 1;
		mat.resize(max_log);
		mat[0] = a;
		for(int j = 1; j < max_log; ++j) {
			mat[j].resize(n - (1 << j) + 1);
			for(int i = 0; i <= n - (1 << j); ++i) {
				mat[j][i] = op(mat[j - 1][i], mat[j - 1][i + (1 << (j - 1))]);
			}
		}
	}

	S prod(int from, int to) const {
		assert(0 <= from && from <= to && to <= n - 1);
		int lg = std::__lg(to - from + 1);
		return op(mat[lg][from], mat[lg][to - (1 << lg) + 1]);
	}

private:
	int n;
	std::vector<std::vector<S>> mat;
};

} // namespace felix
#line 8 "library/tree/hld.hpp"

namespace felix {

struct HLD {
private:
	static constexpr std::pair<int, int> __lca_op(std::pair<int, int> a, std::pair<int, int> b) {
		return std::min(a, b);
	}

public:
	int n;
	std::vector<std::vector<int>> g;
	std::vector<int> subtree_size;
	std::vector<int> parent;
	std::vector<int> depth;
	std::vector<int> top;
	std::vector<int> tour;
	std::vector<int> first_occurrence;
	std::vector<int> id;
	std::vector<std::pair<int, int>> euler_tour;
	sparse_table<std::pair<int, int>, __lca_op> st;

	HLD() : n(0) {}
	explicit HLD(int _n) : n(_n), g(_n), subtree_size(_n), parent(_n), depth(_n), top(_n), first_occurrence(_n), id(_n) {
		tour.reserve(n);
		euler_tour.reserve(2 * n - 1);
	}

	void add_edge(int u, int v) {
		assert(0 <= u && u < n);
		assert(0 <= v && v < n);
		g[u].push_back(v);
		g[v].push_back(u);
	}

	void build(int root = 0) {
		assert(0 <= root && root < n);
		parent[root] = -1;
		top[root] = root;
		dfs_sz(root);
		dfs_link(root);
		st = std::move(sparse_table<std::pair<int, int>, __lca_op>(euler_tour));
	}

	int get_lca(int u, int v) {
		assert(0 <= u && u < n);
		assert(0 <= v && v < n);
		int L = first_occurrence[u];
		int R = first_occurrence[v];
		if(L > R) {
			std::swap(L, R);
		}
		return st.prod(L, R).second;
	}

	bool is_ancestor(int u, int v) {
		assert(0 <= u && u < n);
		assert(0 <= v && v < n);
		return id[u] <= id[v] && id[v] < id[u] + subtree_size[u];
	}

	bool on_path(int a, int x, int b) {
		return (is_ancestor(x, a) || is_ancestor(x, b)) && is_ancestor(get_lca(a, b), x);
	}

	int get_distance(int u, int v) {
		return depth[u] + depth[v] - 2 * depth[(get_lca(u, v))];
	}

	std::pair<int, std::array<int, 2>> get_diameter() const {
		std::pair<int, int> u_max = {-1, -1};
		std::pair<int, int> ux_max = {-1, -1};
		std::pair<int, std::array<int, 2>> uxv_max = {-1, std::array<int, 2>{-1, -1}};
		for(auto [d, u] : euler_tour) {
			u_max = std::max(u_max, std::make_pair(d, u));
			ux_max = std::max(ux_max, std::make_pair(u_max.first - 2 * d, u_max.second));
			uxv_max = std::max(uxv_max, std::make_pair(ux_max.first + d, std::array<int, 2>{ux_max.second, u}));
		}
		return uxv_max;
	}

	int get_kth_ancestor(int u, int k) {
		assert(0 <= u && u < n);
		if(depth[u] < k) {
			return -1;
		}
		int d = depth[u] - k;
		while(depth[top[u]] > d) {
			u = parent[top[u]];
		}
		return tour[id[u] + d - depth[u]];
	}

	int get_kth_node_on_path(int a, int b, int k) {
		int z = get_lca(a, b);
		int fi = depth[a] - depth[z];
		int se = depth[b] - depth[z];
		if(k < 0 || k > fi + se) {
			return -1;
		}
		if(k < fi) {
			return get_kth_ancestor(a, k);
		} else {
			return get_kth_ancestor(b, fi + se - k);
		}
	}

	std::vector<std::pair<int, int>> get_path(int u, int v, bool include_lca) {
		if(u == v && !include_lca) {
			return {};
		}
		std::vector<std::pair<int, int>> lhs, rhs;
		while(top[u] != top[v]) {
			if(depth[top[u]] > depth[top[v]]) {
				lhs.emplace_back(u, top[u]);
				u = parent[top[u]];
			} else {
				rhs.emplace_back(top[v], v);
				v = parent[top[v]];
			}
		}
		if(u != v || include_lca) {
			if(include_lca) {
				lhs.emplace_back(u, v);
			} else {
				int d = std::abs(depth[u] - depth[v]);
				if(depth[u] < depth[v]) {
					rhs.emplace_back(tour[id[v] - d + 1], v);
				} else {
					lhs.emplace_back(u, tour[id[u] - d + 1]);
				}
			}
		}
		std::reverse(rhs.begin(), rhs.end());
		lhs.insert(lhs.end(), rhs.begin(), rhs.end());
		return lhs;
	}

private:
	void dfs_sz(int u) {
		if(parent[u] != -1) {
			g[u].erase(std::find(g[u].begin(), g[u].end(), parent[u]));
		}
		subtree_size[u] = 1;
		for(auto& v : g[u]) {
			parent[v] = u;
			depth[v] = depth[u] + 1;
			dfs_sz(v);
			subtree_size[u] += subtree_size[v];
			if(subtree_size[v] > subtree_size[g[u][0]]) {
				std::swap(v, g[u][0]);
			}
		}
	}

	void dfs_link(int u) {
		first_occurrence[u] = (int) euler_tour.size();
		id[u] = (int) tour.size();
		euler_tour.emplace_back(depth[u], u);
		tour.push_back(u);
		for(auto v : g[u]) {
			top[v] = (v == g[u][0] ? top[u] : v);
			dfs_link(v);
			euler_tour.emplace_back(depth[u], u);
		}
	}
};

} // namespace felix

#line 6 "test/tree/hld/yosupo-Vertex-Add-Path-Sum.test.cpp"
using namespace std;
using namespace felix;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	int n, q;
	cin >> n >> q;
	vector<int> a(n);
	for(int i = 0; i < n; i++) {
		cin >> a[i];
	}
	HLD hld(n);
	for(int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		hld.add_edge(u, v);
	}
	hld.build(0);
	fenwick<long long> fenw(n);
	for(int i = 0; i < n; i++) {
		fenw.add(hld.id[i], a[i]);
		fenw.add(hld.id[i] + hld.subtree_size[i], -a[i]);
	}
	while(q--) {
		int type, u, v;
		cin >> type >> u >> v;
		if(type == 0) {
			fenw.add(hld.id[u], v);
			fenw.add(hld.id[u] + hld.subtree_size[u], -v);
		} else {
			int z = hld.get_lca(u, v);
			long long ans = fenw.get(hld.id[u] + 1) + fenw.get(hld.id[v] + 1) - fenw.get(hld.id[z] + 1);
			if(hld.parent[z] != -1) {
				ans -= fenw.get(hld.id[hld.parent[z]] + 1);
			}
			cout << ans << "\n";
		}
	}
	return 0;
}

Back to top page