[백준 9489] 사촌

문제

증가하는 정수 수열을 이용해서 트리를 만드는 방법은 다음과 같다.

• 첫 번째 정수는 트리의 루트 노드이다.
• 다음에 등장하는 연속된 수의 집합은 루트의 자식을 나타낸다. 이 집합에 포함되는 수의 첫 번째 수는 항상 루트 노드+1보다 크다.
• 그 다음부터는 모든 연속된 수의 집합은 아직 자식이 없는 노드의 자식이 된다. 그러한 노드가 여러 가지 인 경우에는 가장 작은 수를 가지는 노드의 자식이 된다.
• 집합은 수가 연속하지 않는 곳에서 구분된다.

예를 들어, 수열 1 3 4 5 8 9 15 30 31 32를 위의 규칙을 이용해 트리를 만들면 아래 그림과 같이 된다.

두 노드의 부모는 다르지만, 두 부모가 형제(sibling)일 때 두 노드를 사촌이라고 한다.

수열 특정 노드 번호 k가 주어졌을 때, k의 사촌의 수를 구하는 프로그램을 작성하시오.

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다. 각 테스트 케이스의 첫째 줄에는 노드의 수 n과 사촌의 수를 구해야 하는 노드의 번호 k가 주어진다. (1 ≤ n ≤ 1,000, 1 ≤ k ≤ 1,000,000) 다음 줄에는 총 n개의 수가 주어지며, 모든 수는 1보다 크거나 같고, 1,000,000보다 작거나 같다. 입력으로 주어지는 수열은 항상 증가한다. k는 항상 수열에 포함되는 수이다.

입력의 마지막 줄에는 0이 두 개 주어진다.

출력

각 테스트 케이스 마다, k의 사촌의 수를 출력한다.

---

풀이

문제에서 원하는 것을 구하기 전에 먼저 해결해야 하는 부분은 입력값을 조건에 맞게 트리로 만드는 것이다

 

주어진 수는 다음과 같은 조건으로 트리가 형성된다

1. 첫 수는 루트 노드

2. 두번째 수부터 입력되는 연속된 수의 집합은 아직 자식이 없는 노드의 자식이 됨 (오름차순으로)

 

즉, 입력값이 1 3 4 5 8 9 15 30 31 32 라면

1, (3 4 5), (8 9), (15), (30 31 32) 와 같은 연속된 수의 집합 형태로 전처리를 해줘야 한다

 

1은 루트 노드가 되고

3 4 5 는 루트 노드의 자식이 된다

8 9 는 앞서 형성된 3 4 5 중 자식이 없는 가장 작은 수인 3의 자식이 되고

15 는 그 다음 수인 4의 자식

30 31 32는 그 다음 수인 5의 자식이 된다

 

    Q.clear()
    Q.append(nodes[0])
    Tree[nodes[0]] = []
    for idx in range(1,n):
        Tree[nodes[idx]] = []
        if nodes[idx] != nodes[idx-1]+1:	# 연속된 수의 집합에서 벗어나면 부모 변경
            parent = Q.popleft()
        Tree[parent].append(nodes[idx])		# 자식 정보 저장
        parentInfo[nodes[idx]] = parent		# 부모 정보 저장
        Q.append(nodes[idx])

 

입력값을 위와같이 전처리하여 조건에 맞는 트리를 형성하였고 Tree에는 각 노드의 자식 정보, parentInfo에는 각 노드의 부모 정보를 저장하였다.

 

이제 만들어진 트리에서 k의 사촌 수를 구하면 된다

기존에 생각했던 방법으로는 (k가 속한 depth의 노드 수) - (k가 포함된 연속된 수의 집합) 이었다

하지만 이 식이 올바르게 적용되기 위해서는 k의 부모가 해당 depth에서 나머지 모든 노드들과 같은 부모를 갖고 있어야 한다는 오점이 있다

 

따라서 다시 찾아낸 식으로 k의 부모의 부모를 루트 노드로 하는 서브트리를 탐색하였다

이렇게 되면 앞서 말한 오점이 해결된다

 

    # k의 사촌 수 판단
    Q.clear()
    if k == nodes[0]: # k가 루트 노드일 떄 예외
        print(0)
        continue
    elif parentInfo[k] == nodes[0]: # k가 루트 노드의 자식일 때 예외
        print(0)
        continue
    Q.append(parentInfo[parentInfo[k]])
    while len(Q) != 0:
        cycle = len(Q)
        if sequenceK != 0:
            depth = cycle
            break
        for idx in range(cycle):
            currentNode = Q.popleft()
            if len(Tree[currentNode]) != 0:
                for node in Tree[currentNode]:
                    Q.append(node)
                    if node == k:
                        sequenceK = len(Tree[currentNode])

 

k의 부모의 부모 노드부터 BFS로 탐색을 시작했고 (k가 속한 depth의 노드 수) - (k가 포함된 연속된 수의 집합) 로 정답을 구했다

이 과정에서 k가 루트 노드일 때와 k가 루트 노드의 직계자식 인 경우 keyerror가 발생하여 예외처리 해주었다