백준 9489 : 사촌

https://www.acmicpc.net/problem/9489

9489번: 사촌

 

해결 방법 : 트리

• 조건에 따라 트리를 구성하기만 하면 된다.
  • 자식이 최근에 삽입된 노드(cur_p를 index로 갖는 노드)의 마지막 자식 노드 + 1이 현재 number가 아니면, cur_p를 1 증가시켜주기만 하면 된다.
• 사촌 노드는 k의 부모의 부모의 자식들의 자식들이다. 
  • 단, k의 부모는 예외처리 해주어야 할 것.

import sys
input = sys.stdin.readline
from collections import defaultdict

def solution(n, k, nums):
    if k == nums[0]: # root
        return 0

    answer = 0
    T, rT = defaultdict(list), dict() # Tree, reverse_Tree
    cur_p = 0 # nums[p] is current parent node
    T[None], rT[nums[cur_p]] = [nums[cur_p]], None
        
    for i in range(1, n): # Initialize Tree
        if T[nums[cur_p]] != []:
            if nums[i] != T[nums[cur_p]][-1] + 1:
                cur_p += 1
        T[nums[cur_p]].append(nums[i])
        rT[nums[i]] = nums[cur_p]
    
    for p in T[rT[rT[k]]]:
        if p != rT[k]:
            answer += len(T[p])
    
    return answer

if __name__ == '__main__':
    while True:
        n, k = map(int, input().split())
        if n == 0: break
        nums = list(map(int, input().split()))
        result = solution(n, k, nums)
        print(result)