BOJ 21606 - 아침 산책

문제

모든 점이 1(실내), 0(실외)로 구분되는 트리가 주어졌을 때, 산책할 수 있는 경로의 개수를 구하는 문제이다.
경로의 조건은 처음과 끝 지점이 1(실내)이어야 하며, 이외의 점들은 모두 0(실외)여야 한다는 것 뿐이다.
따라서 경로 \(P\)를 뒤집은 경로도 $$P$와 다른 경로로 인정된다.

트리의 정점의 개수 \(N\)은 \(2\le N \le 2\times{10^5}\)를 만족한다.

풀이

서브태스크가 존재하는 문제이다.
입력이 최대 \(2\times{10^5}\)이기 때문에 모든 점들에 대해 그래프 순회로 풀면 2, 5번 서브태스크를 틀리게 된다.

경로의 처음과 끝만 1이고 나머지는 모두 0인 점을 이용해야 한다.
인접해있는 0인 점들을 SCC(강결합요소)처럼 하나의 component로 본다면 이 0 덩어리에 인접한 1의 개수 tcnt를 구할 수 있다. 이때 이 0 덩어리를 통해 나올 수 있는 경로의 개수를 tcnt*(tcnt-1)이 된다.

0인 점들에 대해서 DFS를 수행하면서 0 덩어리, 인접한 1의 개수를 구할 수 있다. 이떄 DFS를 수행하는 반복문에서 0 정점들에 대한 visited 배열(아래 코드에서는 ovis)을 선언해서 이전 정점 순회때 덩어리에 포함되었던 정점에 대해 다시 탐색하지 않게 해야 한다.

소스코드

#include<iostream>
#include<vector>
#include<string>
#include<stack>
#define IN '1'
#define OUT '0'

using namespace std;
using vi=vector<int>;
using lld=long long;

int main(){
	ios::sync_with_stdio(0);
	cin.tie(0);
	int n;
	lld cnt=0;
	string ss;
	cin>>n;
	cin>>ss;
	vi ovis(n, 0);
	vector<vi> adj(n, vi(0));
	for(int i=0;i<n-1;i++){
		int s, e;
		cin>>s>>e;
		adj[--s].push_back(--e);
		adj[e].push_back(s);
	}
	for(int i=0;i<n;i++){
		if(ss[i]==IN) continue;
		if(ovis[i]) continue;
		int tcnt=0;
		vi vis(n, 0);
		stack<int> st;
		ovis[i]=1;
		st.push(i);
		while(!st.empty()){
			int cur=st.top(), f=0;
			vis[cur]=1;
			for(int i=0;i<adj[cur].size();i++){
				int adjv=adj[cur][i];
				if(vis[adjv]==0){
					if(ss[adjv]==IN){
						tcnt++;
						vis[adjv]=1;
					}
					else{
						st.push(adjv);
						ovis[adjv]=1;
						f=1;
						break;
					}
				}
			}
			if(f) continue;
			else st.pop();
		}
		cnt+=tcnt*(tcnt-1);
	}
	for(int i=0;i<n;i++){
		if(ss[i]==OUT) continue;
		for(int j=0;j<adj[i].size();j++){
			if(ss[adj[i][j]]==IN) cnt++;
		}
	}
	cout<<cnt;
	return 0;
}

풀고나서

• 그냥 DFS로 풀면 내 코드는 서브태스크 2도 TLE가 나는데, 다른 사람들의 코드는 2까지 맞는다. 이유는 잘 모르겠다.