백준 15586 : MooTube(Gold)

문제 링크

N <= 100000이고 Q <= 100000일 때, 각 쿼리에서 k와 v가 주어지면 v에서부터 간선의 가중치가 k이상인 간선을 이용하여 몇개의 정점에 도달할 수 있는지 구하는 문제이다.

Brute하게 풀면 쿼리마다 BFS를 이용해주면 되는데, 한번에 O(V+E)만큼 걸리는 BFS의 특성상 10만번의 쿼리를 수행할 수 없다.

따라서 오프라인쿼리와, Union Find를 이용하여 문제를 해결하였다.

유사도 순으로 쿼리를 내림차순으로 정렬하면, 요구하는 유사도가 적어져도 기존에 BFS를 통해 찾았던 간선들은 유효하다. 따라서 그러한 정보를 계속 저장하면서 풀이를 진행하였다.

사용했던 간선을 다시 사용하지 않기 위해 edgidx배열을 이용하였고, 요구하는 유사도가 적어졌을 때, 추가적으로 활성화 할 수 있는 간선을 찾기 위해서 endpoints배열을 이용하였다.

그런데 다른 사람의 코드를 확인해보니 크루스칼 알고리즘을 사용하듯이 간선들을 정렬해준 후 k이상인 간선의 정점들을 merge해준 후, sz[v]의 값을 저장함으로써 문제를 해결하였다. 이때는 간선 정렬에 필요한 O(logn)의 시간복잡도가 소모된다. 결과적으로, endpoints배열을 이용해 시간을 비효율적으로 쓰는 내 풀이보다 훨씬 빨리 해결되었다. 두 풀이 다 제출을 해본 결과, 내 풀이는 376ms, 이 풀이는 112ms로 약 3배 정도 차이가 났다.

문제에서 구하라는 것에만 맹목적으로 매달리지 말고 좀더 크게 보도록 해야겠다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <utility>
#include <string>
#include <queue>

using namespace std;

typedef pair<int,int> pii;

int n,q;

vector<vector<pii>> edge;
vector<pair<pii,int>> query;
vector<int> edgidx;
vector<vector<int>> endpoints;
vector<int> esize;
vector<int> p;
vector<int> sz;

int find(int a)
{
	if(p[a] != a) p[a] = find(p[a]);
	return p[a];
}

bool merge(int u,int v)
{
	u = find(u);
	v = find(v);
	
	if(u == v) return false;
	
	p[u] = v;
	sz[v] += sz[u];
	endpoints[v].insert(endpoints[v].end(),endpoints[u].begin(),endpoints[u].end());
	endpoints[u].clear();
	
	return true;
}

int main(int argc, char** argv) {
	scanf("%d %d",&n,&q);
	
	endpoints = vector<vector<int>>(n);
	edge = vector<vector<pii>>(n);
	sz = vector<int>(n,1);
	esize = vector<int>(n,0);
	edgidx = vector<int>(n,0);
		
	for(int i = 0;i<n-1;i++)
	{
		int v1,v2,w;
		scanf("%d %d %d",&v1,&v2,&w);
		
		v1--; v2--;
		
		edge[v1].push_back({w,v2});
		edge[v2].push_back({w,v1});
		
		esize[v1]++;
		esize[v2]++;
	}
	
	for(int i = 0;i<n;i++)
	{
		p.push_back(i);
		sort(edge[i].begin(),edge[i].end(),greater<pii>());
	}
	
	for(int i = 0;i<q;i++)
	{
		int k,v;
		scanf("%d %d",&k,&v);
		v--;
		
		query.push_back({make_pair(k,v),i});
	}
	
	sort(query.begin(),query.end(),greater<pair<pii,int>>());
	vector<int> ans(q);
	
	for(int i = 0;i<q;i++)
	{
		queue<int> q;
		
		int k = query[i].first.first;
		int fi = find(query[i].first.second);
		q.push(query[i].first.second);
		
		for(auto elem : endpoints[fi])
		{
			q.push(elem);
		}
		
		vector<int> tempendp;
		endpoints[fi].clear();
		
		while(!q.empty())
		{
			int cur = q.front();
			q.pop();
			
			for(;edgidx[cur] < esize[cur];edgidx[cur]++)
			{
				int next = edge[cur][edgidx[cur]].second;
				int w = edge[cur][edgidx[cur]].first;
				
				if(w < k)
				{
					tempendp.push_back(cur);
					break;
				}
				
				if(merge(cur,next))
				{
					q.push(next);	
					
					fi = find(query[i].first.second);
					for(auto elem : endpoints[fi])
					{
						q.push(elem);
					}
				}
			}
		}
		
		fi = find(query[i].first.second);
		endpoints[fi] = tempendp;
		ans[query[i].second] = sz[fi] -1;
	}	
	
	for(auto elem : ans)
	{
		printf("%d\n",elem);
	}
	return 0;
}