BOJ 5825 Yin and Yang

문제 링크입니다.

문제 요약

트리의 각 간선에 0 또는 1을 부여합니다. 어떤 단순 경로에 포함된 0과 1의 개수가 서로 같은 경우 그 경로를 균형 잡힌 경로라고 합시다. 두 개의 균형 잡힌 경로로 쪼갤 수 있는 단순 경로의 수를 구하세요.

풀이

편의 상 0과 1 대신 -1과 1을 부여하겠습니다. 이제 경로의 가중치 합이 0이면 균형 잡힌 경로네요.

센트로이드 분할을 수행합시다. 각 서브트리의 센트로이드 $c$를 지나는 경로의 수를 구한 뒤 다 더하면 원하는 답을 얻을 수 있습니다.

서브트리에서, 루트의 각 자식 별로 DFS를 수행합니다. 이를 통해,

• $c$에서 끝나는 경로
• $c$를 경유하는 경로

의 수를 각각 셀 수 있어야 합니다. 이를 위해, 각 경로마다 경로의 끝에서 $c$ 쪽으로 올라올 때 시작점과 $c$ 외의 정점에서 누적 가중치가 0인 때가 오는 지를 체크할 수 있어야 합니다. 이는 $c$에서 해당 정점으로 내려가는 중에 만난 최댓값과 최솟값 사이에 현재 가중치가 존재함과 동치입니다. 반대로 말하면, 이번에 도달한 정점에서 누적 가중치가 DFS 상에서 처음 만난 값이면 중간에 그 값을 다시 방문할 수 없기 때문에 균형 잡힌 경로가 둘일 수 없습니다.

DFS를 수행하면서 이전까지의 자식 별 DFS 결과의 누적을 관리하는 식으로 계산하면 두 종류의 경로의 수를 ${\cal O}(sz)$에 해결할 수 있습니다. 센트로이드 분할을 했기 때문에 모든 서브트리에 대한 작업을 수행했을 때의 총 시간 복잡도는 ${\cal O}(N \lg N)$이 됩니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;

vector<pii> e[101010];
bool vis[101010];
int sz[101010], yinyang[2][202020];
ll ans;
vector<pii> upd;

int getSize(int now, int pre){
	sz[now] = 1; for(auto& [nxt, nw] : e[now]) if(nxt != pre && !vis[nxt]) sz[now] += getSize(nxt, now);
	return sz[now];
}

int getCent(int now, int pre, int cap){
	for(auto& [nxt, nw] : e[now]) if(nxt != pre && !vis[nxt] && sz[nxt]*2 > cap) return getCent(nxt, now, cap);
	return now;
}

void dfs(int now, int pre, int W, int pmx, int pmn, bool zero){
	ans += yinyang[1][100000-W];
	if(W == 0 && zero) ans++; if(W == 0) zero = 1;
	if(pmx >= W && W >= pmn) ans += yinyang[0][100000-W];
	upd.emplace_back(pmx >= W && W >= pmn, 100000+W);
	for(auto& [nxt, nw] : e[now]) if(nxt != pre && !vis[nxt]) dfs(nxt, now, W+nw, max(pmx, W), min(pmn, W), zero);
}

void solve(int now){
	vis[now = getCent(now, 0, getSize(now, 0))] = 1;
	for(int i = 100000-sz[now]; i <= 100000+sz[now]; i++) yinyang[0][i] = yinyang[1][i] = 0;
	for(auto& [nxt, w] : e[now]) if(!vis[nxt]){dfs(nxt, now, w, 0, 0, 0); for(auto& [i, j] : upd) yinyang[i][j]++; upd.clear();}
	for(auto& [nxt, w] : e[now]) if(!vis[nxt]){solve(nxt);}
}

int main(){
	cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);
	int n; cin >> n;
	for(int i = 1; i < n; i++){
		int u, v, w; cin >> u >> v >> w;
		if(!w) w = -1; e[u].push_back({v, w}); e[v].push_back({u, w});
	}
	solve(1); cout << ans;
}