SWEA 5174 subtree (python)

5174. [파이썬 S/W 문제해결 기본] 8일차 - subtree

문제

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

주어지는 트리는 부모와 자식 노드 번호 사이에 특별한 규칙이 없고, 부모가 없는 노드가 전체의 루트 노드가 된다.

이런 경우의 트리는 부모 노드를 인덱스로 다음과 같은 방법으로 나타낼 수 있다. 자식 노드가 0인 경우는 노드가 자식이 없는 경우이다.

부모	1	2	3	4	5	6
자식1	6	1	0	0	3	4
자식2	0	5	0	0	0	0

[입력]

첫 줄에 테스트케이스의 수 T가 주어진다. 1<=T<=50

다음 줄부터 테스트 케이스의 별로 첫 줄에 간선의 개수 E와 N이 주어지고, 다음 줄에 E개의 부모 자식 노드 번호 쌍이 주어진다.

노드 번호는 1번부터 E+1번까지 존재한다. 1<=E<=1000, 1<=N<=E+1

[출력]

각 줄마다 "#T" (T는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

접근과정

먼저, 문제에 주어진 표를 사용하기 위해 고민을 했다.

부모를 인덱스를 기준으로 하고 2개의 열이 존재하는 array를 만들었다.

문제풀이

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

T = int(input())


def sub_tree(idx):
    # 횟수세기
    global count

    # 좌우 검사
    for i in range(2):
        # 해당되는 곳에 값이 있으면
        if tree[i][idx]:
            # 카운트 증가
            count += 1
            # 그값을 인덱스값으로 활용
            n = tree[i][idx]
            # 다음 세대 확인
            sub_tree(n)


for tc in range(1, T+1):
    E, N = map(int, input().split())
    temp = input().split()

    # 문제에 주어진것 처럼 row 기준으로 0, 1로 활용
    tree = [[0 for _ in range(E+2)]for _ in range(2)]

    for i in range(E):
        # 주어진 인덱스와 값 할당
        idx = int(temp[2*i])
        number = int(temp[2*i+1])
        # 값이 없다면 더해줌
        if tree[0][idx] == 0:
            tree[0][idx] = number
        #  이미 있다면 가지 하나 더 만들어줌
        else:
            tree[1][idx] = number

    # 시작하면서 자동으로 하나 포함
    count = 1
    # N부터 실행
    sub_tree(N)




    print("#{} {}".format(tc, count))