test/ReRooting.test.cpp

• View this file on GitHub
• Last update: 2021-03-26 12:01:06+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/WUPC/3163

Depends on

• Graph/ReRooting.cpp

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/WUPC/3163"
#include "./../Graph/ReRooting.cpp"
#include <iostream>
#include <vector>
using namespace std;
using ll = long long;

vector<ll> a;

struct DP {
	ll sum, val;
	DP() : sum(0), val(0) {}
	DP operator+(const DP& d) const {
		return DP(*this) += d;
	}
	DP& operator+=(const DP& d) {
		sum += d.sum;
		val += d.val;
		return *this;
	}
	DP add_root([[maybe_unused]] int v) const {
		DP res = *this;
		res.val += sum;
		res.sum += a[v];
		return res;
	}
};

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

	int n;
	cin >> n;
	a.assign(n, 0);
	for (ll& i : a) cin >> i;

	ReRooting<DP> dp(n);
	for (int i = 0; i < n - 1; ++i) {
		int u, v;
		cin >> u >> v;
		dp.add_edge(u - 1, v - 1);
	}
	for (auto i : dp.solve()) {
		cout << i.val << endl;
	}
}

#line 1 "test/ReRooting.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/challenges/sources/VPC/WUPC/3163"
#line 2 "Graph/ReRooting.cpp"
#include <vector>

template <class DP> class ReRooting {
	std::size_t n;
	std::vector<std::vector<int>> graph;
	std::vector<std::vector<DP>> dp;
	std::vector<DP> ans;

	DP dfs(int v, int p) {
		DP sum;
		for (std::size_t i = 0; i < graph[v].size(); ++i) {
			int e = graph[v][i];
			DP& dp_e = dp[v][i];
			if (e != p) {
				dp_e = dfs(e, v);
				sum += dp_e;
			}
		}
		return sum.add_root(v);
	}
	void bfs(int v, int p, const DP& dp_par) {
		for (std::size_t i = 0; i < graph[v].size(); ++i) {
			if (graph[v][i] == p) {
				dp[v][i] = dp_par;
			}
		}

		std::vector<DP> dp_left(graph[v].size() + 1);
		for (std::size_t i = 0; i < graph[v].size(); ++i) {
			dp_left[i + 1] = dp_left[i] + dp[v][i];
		}
		std::vector<DP> dp_right(graph[v].size() + 1);
		for (int i = graph[v].size() - 1; i >= 0; --i) {
			dp_right[i] = dp_right[i + 1] + dp[v][i];
		}
		ans[v] = dp_left.back().add_root(v);

		for (std::size_t i = 0; i < graph[v].size(); ++i) {
			int e = graph[v][i];
			if (e != p) {
				bfs(e, v, (dp_left[i] + dp_right[i + 1]).add_root(v));
			}
		}
	}

public:
	ReRooting(std::size_t _n) : n(_n), graph(n), dp(n), ans(n) {}
	ReRooting(const std::vector<std::vector<int>>& _graph)
	    : n(_graph.size()), graph(_graph), dp(n), ans(n) {}
	void add_edge(int v, int u) {
		graph[v].push_back(u);
		graph[u].push_back(v);
	}
	std::vector<DP> solve() {
		for (std::size_t i = 0; i < n; ++i) {
			dp[i].assign(graph[i].size(), DP());
		}
		dfs(0, -1);
		bfs(0, -1, DP());
		return ans;
	}
};

/*
struct DP {
    // edit
    int dp;
    DP(int _dp = 1) : dp(_dp) {}
    DP operator+(const DP& d) const {
        return DP(*this) += d;
    }
    DP& operator+=(const DP& d) {
        // edit
        return *this;
    }
    DP add_root([[maybe_unused]] int v) const {
        DP res = *this;
        // edit
        return res;
    }
};
*/
#line 3 "test/ReRooting.test.cpp"
#include <iostream>
#line 5 "test/ReRooting.test.cpp"
using namespace std;
using ll = long long;

vector<ll> a;

struct DP {
	ll sum, val;
	DP() : sum(0), val(0) {}
	DP operator+(const DP& d) const {
		return DP(*this) += d;
	}
	DP& operator+=(const DP& d) {
		sum += d.sum;
		val += d.val;
		return *this;
	}
	DP add_root([[maybe_unused]] int v) const {
		DP res = *this;
		res.val += sum;
		res.sum += a[v];
		return res;
	}
};

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

	int n;
	cin >> n;
	a.assign(n, 0);
	for (ll& i : a) cin >> i;

	ReRooting<DP> dp(n);
	for (int i = 0; i < n - 1; ++i) {
		int u, v;
		cin >> u >> v;
		dp.add_edge(u - 1, v - 1);
	}
	for (auto i : dp.solve()) {
		cout << i.val << endl;
	}
}

Back to top page