BOJ 31744 Antifreeze

문제 링크입니다.

풀이 역행하기

최종적으로는 임의의 두 난방 기구가 작동하는 교실 사이를 오갈 수 있는지 빠르게 판단할 수 있어야 합니다. 쿼리 수가 아주 작은 것도 아니니, 오프라인 쿼리보다는 그냥 진짜로 모든 쌍에 대해서 미리 구해두는 게 나을 것 같네요.

난방 기구가 있는 세 교실 A, B, C가 있다고 합시다. A에서 B로 갈 수 있고, B에서 C로 갈 수 있다면 A에서 C로도 갈 수 있어요. A, B, C가 하나의 세트로 묶였다고 생각할 수도 있겠네요. 다른 난방 기구가 있는 교실 D에서 A, B, C 중 한 곳이라도 오갈 수 있으면 나머지 두 교실도 오갈 수 있는 게 됩니다.

그러니 어떻게 잘 처리해서 서로 왔다갔다할 수 있는 난방 기구가 있는 교실끼리를 하나의 세트로 묶인 상태로 만들어 줍시다. 분리 집합을 활용하면 전처리 후 쿼리 당 $\mathcal{O}(\alpha(N))$ 정도에 처리 가능합니다. 그냥 대충 쿼리 당 $\mathcal{O}(1)$ 이라고 생각해도 문제 없습니다.

난방 기구가 없는 교실은 집합으로 묶지 않습니다. 그 교실에 도달했을 때의 온도를 고려하면서 묶어야 하는데, 그러면 경우의 수가 너무 많아집니다.

트리 성질 활용하기

주어지는 그래프는 트리 구조를 이룹니다. 임의의 두 정점을 잇는 단순 경로가 유일하게 존재합니다.

적당한 정점을 트리의 루트로 잡았다고 합시다. 어떤 두 정점 사이를 오가는 경로는 반드시 두 정점의 최소 공통 조상을 지납니다. 이는 즉슨 공통 조상까지 올라갔을 때 그 조상으로부터 두 정점까지의 거리 합이 T를 초과한다면 합쳐질 일이 없다는 의미입니다.

트리에서 DFS를 돌면서, 각 정점마다 그 정점에 닿는 서로 오갈 수 있음이 확인된 난방 기구가 있는 정점 세트들까지의 거리들을 관리하고 있다고 해 봅시다. 그러면, 그중 가장 가까운 것과 다른 모든 정점 세트의 거리 합을 확인해 합쳐주는 식으로 조상을 거쳐 만나는 정점들을 체크해 줄 수 있습니다. 현재 확인 중인 정점에 난방 기구가 있다면 그럴 필요 없이 그냥 바로 현재 정점과 연결해 주면 충분합니다.

경로 관리하기

해당 기능을 나이브하게 구현하면 정점별로 경로를 계속 옮겨 저장해 주어야 하니 $\mathcal{O}(N^2)$으로 시간 초과를 받습니다. 여기서 Smaller-to-Larger 테크닉을 사용합니다. 경로를 std::set과 같이 정렬성을 유지하는 자료 구조로 관리하되, 부모 정점에서 자식 정점이 가지고 있던 경로들을 가져올 때 자식 정점이 가지고 있는 경로 수가 더 많으면 두 set을 서로 바꿔준 뒤 원래 가지고 있던 경로들을 다시 불러오는 방식을 사용합니다. 그러면, 각 원소마다 한 번 위치를 옮겼을 때 집합의 크기가 2배 이상이 되어, 최대 $\mathcal{O}(\log N)$번만 이동함을 알 수 있습니다.

이상의 내용을 구현하면 최종 시간 복잡도 $\mathcal{O}(N \log^2 N)$에 문제를 해결할 수 있습니다.

코드

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

int n, t, par[101010], offset[101010];
bool heat[101010];
set<pii> st[101010];
vector<pii> e[101010];

int f(int x){
	return par[x] < 0 ? x : par[x] = f(par[x]);
}

void u(int x, int y){
	x = f(x); y = f(y);
	if(x == y) return;
	if(par[x] > par[y]) swap(x, y);
	par[x] += par[y]; par[y] = x;
}

void dfs(int now, int pre){
	for(auto [nxt, w] : e[now]){
		if(nxt == pre) continue;
		dfs(nxt, now);
		offset[nxt] += w;
		if(heat[now]){
			for(auto [nw, idx] : st[nxt]){
				if(nw+offset[nxt] > t) break;
				u(now, idx);
			}
		}
		else{
			if(st[nxt].size() > st[now].size()){
				swap(st[nxt], st[now]);
				swap(offset[nxt], offset[now]);
			}
			for(auto [nw, idx] : st[nxt]){
				if(nw+offset[nxt] > t) break;
				st[now].insert({nw+offset[nxt]-offset[now], idx});
			}
		}
	}
	if(heat[now]) return;
	auto wi = *st[now].begin();
	vector<pii> bye;
	for(auto [nw, idx] : st[now]){
		if(wi.second == idx) continue;
		if(wi.first+nw+2*offset[now] > t) break;
		u(wi.second, idx);
		bye.push_back({nw, idx});
	}
	for(auto p : bye) st[now].erase(p);
	bye.clear();
}

int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(false);
	memset(par, -1, sizeof(par));
	cin >> n >> t;
	for(int i = 1; i <= n; i++){
		cin >> heat[i];
		if(heat[i]) st[i].insert({0, i});
	}
	for(int i = 1; i < n; i++){
		int x, y, w; cin >> x >> y >> w;
		e[x].push_back({y, w});
		e[y].push_back({x, w});
	}
	dfs(1, 0);
	int Q; cin >> Q; while(Q--){
		int x, y; cin >> x >> y;
		cout << (f(x) == f(y)) << '\n';
	}
}