[백준][Python] 9489번 사촌

9489번: 사촌

 

코드

import sys
while 1:
    n,k = map(int,sys.stdin.readline().split())
    if n == 0 and k == 0:
        break
    else:
        dic = {}
        n_list = list(map(int,sys.stdin.readline().split()))
        cnt = 0
        
        parent = []
        if len(n_list) > 1:
            next = n_list[1]
        parent.append(next)
        for i in range(2,n):
            if (n_list[i]-1) == next:
                parent.append(n_list[i])
                next = n_list[i]
                
            else:
                dic[n_list[cnt]] = parent
                parent = []
                parent.append(n_list[i])
                next = n_list[i]
                cnt+=1
            if i == n-1:
                dic[n_list[cnt]] = parent   
    
        for i in dic.keys():  # 부모 노드 찾기
            if k in dic[i]:
                root = i
                break
        found = 0
       
        for i in dic.keys():
            if root in dic[i]:
                par = i
                found = 1
                break
        if found == 0:
            print(0)
        else:
            ans = 0
            for i in dic[par]:
                if i != root and (i in dic):
                    ans += len(dic[i])
            print(ans)

이 코드에서는 루트 노드를 기준으로 모든 자식 노드를 딕셔너리에 저장했다. 다음에 오는 수가 현재 수와 1 차이만 나면 parent 리스트에 계속 append 해주고, 그렇지 않으면 딕셔너리에 지금까지의 parent 리스트를 저장한 후 새로운 리스트를 작성해주었다. cnt 값도 이때 1씩 증가시켜서 부모노드 번호를 알아낸다. 

주어진 노드 번호 k와 부모는 다르지만 부모의 부모가 같은 노드를 찾고, 그 번호가 딕셔너리에서 value로 갖는 리스트의 길이를 모두 더해주면 답이 나온다. 

답은 맞게 나왔는데 너무 코드가 쓸데없이 길어진 거 같아서 다른 사람들 코드를 찾아봐서 다시 풀어봤다. 

import sys
from collections import defaultdict
while 1:
    n,k = map(int,sys.stdin.readline().split())
    if n == 0 and k == 0:
        break
    else:
        dic = defaultdict(int)
        n_list = list(map(int,sys.stdin.readline().split()))
        cnt = 0
        for i in range(1,n):
            dic[n_list[i]] = n_list[cnt]
            if i<n-1 and n_list[i]+1 < n_list[i+1]:
                cnt += 1  # 부모 다름
        ans = 0
        if dic[dic[k]]:  # 부모의 부모노드가 존재하면
            for j in n_list:
                if dic[dic[j]] == dic[dic[k]] and dic[j] != dic[k]:
                    ans+=1

        print(ans)

딕셔너리에 들어있지 않은 key값에 접근하려 하면 오류가 떠서 defaultdict를 사용했다. 이러면 루트 노드와 같이 부모노드가 존재하지 않아 딕셔너리에 넣지 않은 것들에도 접근이 가능해진다. 첫번째 코드에서는 이걸 몰라서 더 복잡하게 코딩을 한 것 같다. if dic[dic[k]] 와 같은 if문도 이래서 사용 가능했다. 

나머지는 위와 같이 부모의 부모가 같고, 부모노드는 서로 다른 노드의 수를 구해서 출력해줬다.