[Atcoder] ABC303 :: E - A Gift From the Stars

E - A Gift From the Stars

문제 해석이 좀 어렵다.
트리가 주어지고, 여기서 star를 찾고자 하는 문제이다. 각 노드는 단 하나의 star에만 속할 수 있고, K - star는 star node와 연결된 노드의 개수로 결정이 된다. 그리고 Tree 형태로 주어지기 때문에 모든 노드는 star에 속할 수 밖에 없다.
문제를 위와같이 이해하면 DFS를 떠올려 볼 수 있다. 그리고 star 노드와 연결된 노드의 개수는 미리 O(N)으로 전처리 해놓을 수 있으므로 이 정보를 활용해보기로 하자.

이제 star 노드를 찾아야 하는데 리프 노드를 하나 찾고, 해당 노드 부터 DFS 탐색을 진행하면 쉽게 구할 수 있다.
만약 현재 노드가 star 노드라면 현재 노드와 연결된 노드와 연결된 노드는 무조건 star 노드가 되므로 이 노드들을 전부 찾고 연결된 정보를 활용하면 해결할 수 있다.

그래서 방문을 표시하기 위한 배열을 하나 선언해 놓고, star의 꼭짓점은 0으로, 중심은 1로, 그다음 꼭짓점은 2로 표시를 수행하는 방식으로 탐색을 진행했고 우리가 찾고자 하는 노드는 '1'로 표시된 노드이므로 이 노드와 연결된 개수들을 전부 구한 후 출력하면 된다.

#include<iostream>
#include <vector>
#include <cstring>
#include <algorithm>
#define NUM 200010
using namespace std;
int N, visited[NUM], cnt[NUM], start;
vector<int> adj[NUM];
vector<int> answer;

void dfs(int curr, int pos) {
    if (pos == 1) {
        answer.push_back(cnt[curr]);
    }
    for (auto next: adj[curr]) {
        if (visited[next] != -1) continue;
        visited[next] = (pos + 1) % 3;
        dfs(next, (pos + 1) % 3);
    }
}

int main()
{
    memset(visited, -1, sizeof(visited));
    memset(cnt, 0, sizeof(cnt));
    
    cin >> N;
    
    for (int i=1; i<N; i++) {
        int x, y;
        cin >> x >> y;
        adj[x].push_back(y);
        adj[y].push_back(x);
        cnt[x]++;
        cnt[y]++;
    }
    
    for (int i=1; i<=N; i++) {
        if (cnt[i] == 1) {
            start = i;
            break;
        }
    }
    
    visited[start] = 0;
    dfs(start, 0);
    
    sort(answer.begin(), answer.end());
    for (auto ans: answer) {
        cout << ans << ' ';
    }
}​