[SWEA 5174] subtree python

문제

트리의 일부를 서브 트리라고 한다. 주어진 이진 트리에서 노드 N을 루트로 하는 서브 트리에 속한 노드의 개수를 알아내는 프로그램을 만드시오.

 

풀이

import sys
sys.stdin = open("input.txt", "r")

def subtree(n):
    global cnt
    cnt += 1
    for i in range(2):
      if arr[i][n]:
          subtree(arr[i][n])

T = int(input())
for test_case in range(1, T + 1):
    e, n = map(int, input().split())
    arr = [[0 for _ in range(e+2)]for _ in range(2)]
    inputList = list(map(int, input().split()))

    for i in range(0, len(inputList), 2):
      parent, sibiling = inputList[i], inputList[i+1]
      if arr[0][parent]:
          arr[1][parent] = sibiling
      else:
          arr[0][parent] = sibiling
    
    cnt = 0
    subtree(n)
    print("#%d %d" % (test_case, cnt))

• e(간선)과 n(노드)의 개수를 입력받습니다. *트리에서 (간선의 개수+1)은 (노드의 개수)와 같습니다.
• arr 에서 인덱스는 부모의 노드 번호로 생각하고, arr[0]은 첫 번째 자식, arr[1]은 두 번째 자식으로 생각합니다.
• input에서 부모 자식 순이 반복되기 때문에, 2씩 건너 뛰며 반복문을 돌며 처리합니다.