[백준 9489번] 사촌 (python)

문제

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

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

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

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

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

 

기록

(트리 4일차) sibling 관계를 잘 이해못해서 헤맸다. 앞으로는 문제를 꼼꼼히 읽으려고 노력해야겠다. 재귀나 dfs bfs 안쓰고 위로 아래로 두 번씩만(조상-부모-자식) 탐색하면 되는거라 코드가 어렵진 않았다.

 

코드

from sys import stdin
from collections import defaultdict

while True:

    # 변수 초기화
    n, k = map(int, stdin.readline().split())
    if n==0 : break
    seq = list(map(int, stdin.readline().split()))
    tree = defaultdict(list)
    rootidx = -1
    for i in range(1, len(seq)):
        if seq[i]-seq[i-1]>1 :
            rootidx+=1
        tree[seq[rootidx]].append(seq[i])
        tree[seq[i]].append(seq[rootidx])

    # k의 부모, 조상 노드 찾기
    if rootidx==-1 :
        print(0)
        continue
    k_par = tree[k][0]
    if k_par > k :
        print(0)
        continue
    k_anc = tree[k_par][0]
    if k_anc > k_par :
        print(0)
        continue

    # k 부모를 제외한 손자들 계산
    sib = 0
    for par in tree[k_anc]:
        if par<k_anc or par==k_par : continue
        sib+=len(tree[par][1:])

    print(sib)

1. edge 정보는 수열 규칙에 맞게 처리해서 양방향으로 모두 저장했다. 나중에 루트랑 자식 edge 구분할 때는 자신보다 숫자가 작은지 큰 지 구분해서 갈 수 있기 때문에 따로 구분안했다.

 

2. k부터 트리따라 올라가면서 부모노드와 조상노드 찾아서 변수에 저장

 

3. k의 조상노드에서 트리 아래로 내려가면서 k의 부모노드를 제외한 다른 부모노드가 갖고있는 자식노드의 개수 세서 출력하면 끝!