CF1000G. Two-Paths(树形dp)

CF1000G. Two-Paths

Solution

我们发现除了树上$(x,y)$最短路径上的边经过一次，其余边要么走$0$次，要么走$2$次。

因此考虑先假设每条边走两次，最后把走一次的边的贡献加上。

我们把从$(x,y)$路径上的点扩展出去连痛块的贡献拆成三部分：

1. $x,y$的子树中的部分。
2. $x,y$路径上不是端点的点延伸出去且不经过$(x,y)$路径的部分。
3. $x,y$的$lca$向上延伸的部分。

我们可以通过树形$dp$来求出这三种贡献，最后询问就形如路径求和，倍增或者树剖即可。

时间复杂度$O(n\log n)$。

Code

#include <bits/stdc++.h>

using namespace std;
 
template<typename T> inline bool upmin(T &x, T y) { return y <= x ? x = y, 1 : 0; }
template<typename T> inline bool upmax(T &x, T y) { return x <= y ? x = y, 1 : 0; }
 
#define MP(A,B) make_pair(A,B)
#define PB(A) push_back(A)
#define SIZE(A) ((int)A.size())
#define LEN(A) ((int)A.length())
#define FOR(i,a,b) for(int i=(a);i<(b);++i)
#define fi first
#define se second
 
typedef long long ll;
typedef unsigned long long ull;
typedef long double lod;
typedef pair<int, int> PR;
typedef vector<int> VI; 
 
const lod eps = 1e-9;
const lod pi = acos(-1);
const int oo = 1 << 30;
const ll loo = (1ll << 62) - 1;
const int MAXN = 600005;
const int mods = 998244353;
const int MX = 100000;
const int inv2 = (mods + 1) >> 1;
const int INF = 0x3f3f3f3f; //1061109567
/*--------------------------------------------------------------------*/
 
namespace FastIO{
	constexpr int SIZE = (1 << 21) + 1;
	int num = 0, f;
	char ibuf[SIZE], obuf[SIZE], que[65], *iS, *iT, *oS = obuf, *oT = obuf + SIZE - 1, c;
	#define gc() (iS == iT ? (iT = ((iS = ibuf) + fread(ibuf, 1, SIZE, stdin)), (iS == iT ? EOF : *iS ++)) : *iS ++)
	inline void flush() {
		fwrite(obuf, 1, oS - obuf, stdout);
		oS = obuf;
	}
	inline void putc(char c) {
		*oS ++ = c;
		if (oS == oT) flush();
	}
	inline void getc(char &c) {
		for (c = gc(); !isalpha(c) && c != EOF; c = gc());
	}
	inline void reads(char *st) {
		char c;
		int n = 0;
		getc(st[++ n]);
		for (c = gc(); isalpha(c) ; c = gc()) st[++ n] = c;
		st[++ n] = '\0';
	}
	template<class I>
	inline void read(I &x) {
		for (f = 1, c = gc(); c < '0' || c > '9' ; c = gc()) if (c == '-') f = -1;
		for (x = 0; c >= '0' && c <= '9' ; c = gc()) x = (x << 3) + (x << 1) + (c & 15);
		x *= f;
	}
	template<class I>
	inline void print(I x) {
		if (x < 0) putc('-'), x = -x;
		if (!x) putc('0');
		while (x) que[++ num] = x % 10 + 48, x /= 10;
		while (num) putc(que[num --]);
	}
	inline void putstr(string st) {
		for (int i = 0; i < (int)st.size() ; ++ i) putc(st[i]);
	}
	struct Flusher_{~Flusher_(){flush();}} io_Flusher_;
}
using FastIO :: read;
using FastIO :: putc;
using FastIO :: putstr;
using FastIO :: reads;
using FastIO :: print;






struct enode{ int nxt, to, c; } e[MAXN << 1];
int head[MAXN], C[MAXN], fa[MAXN][20], dep[MAXN], a[MAXN], Log[MAXN], edgenum = 0;
ll f[MAXN], g[MAXN], h[MAXN], sum[MAXN][20];

void add(int u, int v, int c) {
	e[++ edgenum] = (enode){head[u], v, c}, head[u] = edgenum;
}
void tree_dp(int x, int father) {
	for (int i = head[x]; i ; i = e[i].nxt) {
		int v = e[i].to;
		if (v == father) continue;
		tree_dp(v, x);
	}
	f[x] = a[x];
	for (int i = head[x]; i ; i = e[i].nxt) {
		int v = e[i].to, c = e[i].c;
		if (v == father) continue;
		f[x] += max(f[v] - c - c, 0ll);
	}
	for (int i = head[x]; i ; i = e[i].nxt) {
		int v = e[i].to, c = e[i].c;
		if (v == father) continue;
		g[v] = f[x] - max(f[v] - c - c, 0ll);
	}
}
void dfs(int x, int father) {
	dep[x] = dep[father] + 1, fa[x][0] = father, sum[x][0] = g[x] - C[x];
	for (int i = 1; i <= Log[dep[x]] ; ++ i) fa[x][i] = fa[fa[x][i - 1]][i - 1], sum[x][i] = sum[x][i - 1] + sum[fa[x][i - 1]][i - 1];
	for (int i = head[x]; i ; i = e[i].nxt) {
		int v = e[i].to, c = e[i].c;
		if (v == father) continue;
		h[v] = max(h[x] + f[x] - max(f[v] - c - c, 0ll) - c - c, 0ll);
		C[v] = c;
		dfs(v, x);
	}
}
int getlca(int x, int y) {
	if (dep[x] < dep[y]) swap(x, y);
	for (int i = Log[dep[x]]; i >= 0 ; -- i)
		if (dep[fa[x][i]] >= dep[y]) x = fa[x][i];
	if (x == y) return x;
	for (int i = Log[dep[x]]; i >= 0 ; -- i)
		if (fa[x][i] != fa[y][i]) x = fa[x][i], y = fa[y][i];
	return fa[x][0];
}
int jump(int u, int d) {
	for (int i = Log[d]; i >= 0 ; -- i)
		if ((d >> i) & 1) u = fa[u][i];
	return u;
}
ll getsum(int u, int d) {
	ll ans = 0;
	for (int i = Log[d]; i >= 0 ; -- i)
		if ((d >> i) & 1) ans += sum[u][i], u = fa[u][i];
	return ans;
}

signed main() {
#ifndef ONLINE_JUDGE
	freopen("a.in", "r", stdin);
#endif	
	int n, Case;
	read(n), read(Case);
	for (int i = 1; i <= n ; ++ i) read(a[i]);
	for (int i = 1, u, v, c; i < n ; ++ i) read(u), read(v), read(c), add(u, v, c), add(v, u, c);
	tree_dp(1, 0);
	dep[0] = -1, Log[1] = 0;
	for (int i = 2; i <= n ; ++ i) Log[i] = Log[i >> 1] + 1;
	dfs(1, 0);

//	for (int i = 1; i <= n ; ++ i) cout << i << ":" << f[i] << " " << g[i] << " " << h[i] << endl;
	while (Case --) {
		int u, v;
		read(u), read(v);
		if (dep[u] < dep[v]) swap(u, v);

		ll ans = 0;
		int t = getlca(u, v);
		if (u == v) ans = h[u] + f[u];
		else if (v == t) {
			int uu = jump(u, dep[u] - dep[t] - 1);
			ans = h[t] + f[t] - max(f[uu] - C[uu] - C[uu], 0ll) + getsum(u, dep[u] - dep[t] - 1) + f[u] - C[uu];
		}
		else {
			int uu = jump(u, dep[u] - dep[t] - 1), vv = jump(v, dep[v] - dep[t] - 1);
			ans = h[t] + f[t] - max(f[uu] - C[uu] - C[uu], 0ll) - max(f[vv] - C[vv] - C[vv], 0ll) - C[uu] - C[vv];
			ans += getsum(u, dep[u] - dep[t] - 1) + getsum(v, dep[v] - dep[t] - 1) + f[u] + f[v]; 	
		}
		print(ans), putc('\n');
	}
	return 0;
}