[파이썬] 5174. subtree

※ SW Expert 아카데미의 문제를 무단 복제하는 것을 금지합니다.


트리의 일부를 서브 트리라고 한다. 주어진 이진 트리에서 노드 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는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

입력	출력
3 5 1 2 1 2 5 1 6 5 3 6 4 5 1 2 6 6 4 6 5 4 1 5 3 10 5 7 6 7 4 6 9 4 11 9 5 11 8 5 3 5 2 8 1 8 10	#1 3 #2 1 #3 3

 

🎀 나의 풀이

 

# 실패한 코드
T = int(input())


# 여러개의 테스트 케이스가 주어지므로, 각각을 처리합니다.

def search(N):
    global result
    result += 1
    if tree_node[N]:
        search(tree_node[N])
        #search(tree_node.pop(N))
    return result


for test_case in range(1, T + 1):
    E, N = map(int, input().split())  # 간선의 개수
    tree = list(map(int, input().split()))
    tree_node = [[] for _ in range(E + 2)]  # 0 제외
    idx = 1
    for i in tree[::2]:
        tree_node[i].append(tree[idx])
        idx += 2
    print(tree_node)
    result = 0
    print(search(N))

 

TypeError : list indices must be integers or slices, not list

재귀호출한 후 인덱스 값으로 int가 아닌 리스트가 들어가니까(tree_node[N]값은 리스트 타입) 에러가 뜬다.

두번째 문제는 노드가 2개일 경우를 고려안함.

아래는 문제를 성공적으로 해결한 코드

 

T = int(input())
# 여러개의 테스트 케이스가 주어지므로, 각각을 처리합니다.

def search(N):
    global result
    result += 1
    for i in tree_node[N]:		# 자식이 둘일 수도 있음.
        search(i)

for test_case in range(1, T + 1):
    E, N = map(int, input().split())  # 간선의 개수
    tree = list(map(int, input().split()))		# 노드들
    tree_node = [[] for _ in range(E + 2)]  # 0 제외
    idx = 1
    for i in tree[::2]:
        tree_node[i].append(tree[idx])		# 부모 노드 = 인덱스, 값 = 부모 노드의 자식
        idx += 2
    result = 0
    search(N)
    print(f"#{test_case} {result}")

 

부모 노드를 인덱스로 하는 부모-자식간의 연결 정보를 나타내는 tree_node리스트를 만든다.

노드의 수는 간선의 수+1이다.

노드가 1부터 시작이기 때문에 간선의 수E+2를 했다.

재귀호출을 이용해서 서브트리에 속한 노드의 개수를 센다.

 

 

🎀 다른사람 풀이

 

def dfs(idx):
    global count
    # 순회를 할때마다 카운트
    count+=1
    # 자식노드를 순회
    for i in Tree[idx] :
        dfs(i)

for t in range(int(input())) :
    # 노드의 개수와 출력할 노드
    E, N=map(int, input().split())
    temp_list=input().split()
    # 노드의 개수만큼 리스트 만들기
    Tree=[[] for _ in range(E+2)]
    for idx, i in enumerate(range(0, E*2, 2)):
        a=int(temp_list[i])
        b=int(temp_list[i+1])
        #부모노드를 찾아서 자식노드 표현
        Tree[a].append(b)
    count=0
    dfs(N)
    #결과값 출력
    print("#{} {}".format(t+1, count))

 

나랑 비슷한 알고리즘.

 

def inorder(node):
  global cnt
  if node == 0:
    return
  # 전위
  cnt+=1
  inorder(left[node])
  # 중위
  inorder(right[node])
  # 후위

for test_case in range(1, int(input())+1) :
  E, N=map(int, input().split())
  arr=list(map(int, input().split()))
  # 구현1
  left=[0]*(E+2)  # 첫번째 자식
  right=[0]*(E+2) # 두번째 자식
  # 구현2
  for i in range(0, len(arr), 2):
    papa, baby=arr[i], arr[i+1]   # 부모 - 자식
    if left[papa] : # 0이 아니고 첫번째에 값이 있으면
      right[papa]=baby    # 두번째에 보관
    else :  # 0이면
      left[papa]=baby     # 첫번째에 보관

  cnt=0
  inorder(N)
  print('#{} {}'.format(test_case, cnt))

 

첫번째 자식 left, 두번째 자식 right로 나눴다.

왼쪽에 값이 있으면 오른쪽에 저장한다.

subtree를 찾으려는 node를 넣고 inorder함수 호출.

왼쪽, 오른쪽 탐색한다.

node가 0이면 return한다.