본문 바로가기

코딩테스트

[Python] 트리 문제 풀이 정리



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

  • 완전 이진 트리 + 중위순회 문제
  • 중위순회의 결과를 보고 트리를 만들 수 있어야 한다. 손으로는 풀 수 있는데...코드로는 구현을 못하겠다..
# 정답 풀이
import sys
input = sys.stdin.readline

K = int(input())
tree = [[] for _ in range(K)]
p = list(map(int, input().split()))

def inOrder(p, depth):
    # 트리의 root index 찾기
    mid = (len(p)//2)

    # 해당 깊이에 해당하는 node 추가
    tree[depth].append(p[mid])

    if len(p)==1:
        return
    inOrder(p[:mid], depth+1)
    inOrder(p[mid+1:], depth+1)

inOrder(p, 0)

for i in tree:
    print(*i)
  • 중위순회 입력 리스트(p)에서 가운데 값은 항상 root index가 된다. 
  • 입력받은 트리 혹은 서브트리의 가운데 값을 찾고, 그때 트리의 level(depth)에 맞게 tree 리스트에 가운데 node를 추가 한다.
  • 만약 len(p)가 1이라면, 즉 해당 노드가 리프 노드라면 return하고, 아니라면 재귀 형식으로(DFS) 왼쪽과 오른쪽 순서대로 서브트리를 탐색해준다. 서브트리이므로 depth도 +1 해주어야 함.

 

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

  • 리프 노드가 홀수인경우 성원이가 이길 수 있다.(Yes)
  • 따라서 탐색을 수행하면서, 리프노드의 개수를 센다. 부모-자식 관계가 확실하게 주어지지 않으므로...리프 노드임을 어떻게 확신해야 할까?
# 처음 풀이(틀림)
import sys
sys.setrecursionlimit(500000)
input = sys.stdin.readline

K = int(input())
tree = [[] for _ in range(K+1)]
visited = [False]*(K+1)
leaf_node = 0

for _ in range(K-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)

def DFS(num):
    global leaf_node
    visited[num]=True
    for i in tree[num]:
        if not visited[i]:
            if len(tree[i])==1 and visited[tree[i][0]]==True:
                leaf_node+=1
            DFS(i)

DFS(1)
if leaf_node%2==0:
    print("No")
else:
    print("Yes")
  • tree리스트의 길이가 1이고, 연결된 노드가 이미 방문한 노드라면 해당 노드는 리프노드라고 가정한다.
  • 시간초과
# 정답 풀이
import sys
sys.setrecursionlimit(500000)
input = sys.stdin.readline

K = int(input())
tree = [[] for _ in range(K+1)]

for _ in range(K-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)

stack = [[1, 0]]  # stack을 이용한 DFS구현
chk = [0] * (K+1)
chk[1] = -1
ls = []

while stack:
    cn, l = stack.pop()
    chk[cn] = 1

    if cn != 1 and len(tree[cn]) == 1:
        ls.append(l)
        continue

    for nn in tree[cn]:
        if chk[nn] == 0:
            stack.append([nn, l+1])

leaf_node = sum_ls = sum(ls)

if leaf_node%2==0:
    print("No")
else:
    print("Yes")
  • 탐색을 DFS가 아닌 stack을 이용하여 구현한다.
  • stack[n][0]==cn 현재 노드 번호, stack[n][1]==l은 현재 노드 n의 level(depth)를 뜻한다. 
  • 현재 노드를 방문하지 않았고, 연결된 노드가 하나뿐인 경우(len(tree[cn])==1) 리프 노드 이므로 ls에 넣는다.
  • visited 여부는 chk리스트로 판별한다.

 

 

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

  • 친구의 친구, 즉 한 다리 까지만 탐색한다.
  • 또한 주인공 상근이의 친구만 탐색하면 되므로 시작 탐색 노드 friends[1]과 연결된 친구를 탐색한다.
import sys
from collections import deque
input = sys.stdin.readline

N = int(input())
M = int(input())
friends = [[] for _ in range(N+1)]
visited = [False]*(N+1)

for _ in range(M):
    a, b = map(int, input().split())
    friends[a].append(b)
    friends[b].append(a)
# print(friends)

ans = 0
visited[1]=True
for i in friends[1]:
    if not visited[i]:
        visited[i]=True
        ans +=1
    for j in friends[i]:
        if not visited[j]:
            visited[j]=True
            ans +=1

print(ans)
  • 만약 visited[i], 탐색을 마친 상근이의 친구이더라도 친구의 친구는 항상 탐색을 해주어야 한다.

 

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

# 내 풀이
import sys
from collections import deque
input = sys.stdin.readline

N = int(input())
A = list(map(int, input().split()))
visited = [-1]*(N+1)
depth = [0]*(N+1)

def BFS(v):
    queue = deque()
    queue.append(v)

    while queue:
        c_v = queue.popleft()
        
        for i in range(3):
            n_v = A[c_v]+i
            # print(n_v, queue, visited)
            if n_v<N+1 and visited[n_v]==-1 :
                if A[n_v]+i>=N:
                    visited[n_v]=visited[c_v]+1
                    return
                else:
                    visited[n_v]=visited[c_v]+1
                    queue.append(n_v)
visited[0]=1
BFS(0)
print(visited)
  • BFS로 풀려다 실패
# 정답 풀이
from collections import deque
import sys
input = sys.stdin.readline

# 입력
N = int(input())
graph = list(map(int, input().split()))

if N == 1: # N이 1칸만 있는 경우 점프 X
    print(0)
else:
    # bfs
    visited = [0]*(N+1)
    queue = deque([1])
    
    while queue:
        x = queue.popleft()
        if x + graph[x-1] >= N:
            print(visited[x]+1)
            break
        for i in range(1, graph[x-1]+1):
            nx = x+i
            if visited[nx] == 0:
                queue.append(nx)
                visited[nx] = visited[x]+1
    else:
        print(-1)
  • Ai이하만큼 떨어진 곳으로 이동 가능..문제를 제대로 이해하지 못하고 풀었다.
# 정답 풀이
import sys

n = int(input())
jump = list(map(int, sys.stdin.readline().split()))
dp = [n + 1] * n
dp[0] = 0

for i in range(n):
	for j in range(1, jump[i] + 1):
		if i + j < n:
			dp[i + j] = min(dp[i] + 1, dp[i + j])

if dp[n - 1] == n + 1:
	print(-1)
else:
	print(dp[n - 1])
  • 탐색보다는 DP로 푸는게 더 정석인 문제. 생각을 못했다.

 

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

from collections import deque
import sys
sys.setrecursionlimit(100000)
input = sys.stdin.readline

# 입력
N, M = map(int, input().split())
visited_W = [[False]*N for _ in range(M)]
visited_B = [[False]*N for _ in range(M)]
S = [] # 전쟁터
for _ in range(M):
    S.append(input().strip())

dx = [-1,1,0,0]
dy = [0,0,-1,1]

def DFS_W(x, y, n): # 내 병사
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]
        if (0<=nx<N) and (0<=ny<M) and not visited_W[ny][nx] and S[ny][nx]=='W':
        # 다음에 이동할 장소가 유효하고 내 병사일 때
            visited_W[ny][nx]=True
            n=DFS_W(nx,ny, n+1)
    return n

def DFS_B(x, y, n): # 상대 병사
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]
        if (0<=nx<N) and (0<=ny<M) and not visited_B[ny][nx] and S[ny][nx]=='B':
        # 다음에 이동할 장소가 유효하고 내 병사일 때
            visited_B[ny][nx]=True
            n=DFS_B(nx,ny, n+1)
    return n

ans_W = 0
ans_B = 0
for i in range(M):
    for j in range(N):
        if not visited_W[i][j] and S[i][j]=='W': # 방문하지 않았고 내 병사면
            visited_W[i][j] = True
            ans_W += DFS_W(j,i, 1)**2
        elif not visited_B[i][j] and S[i][j]=='B':
            visited_B[i][j] = True
            ans_B += DFS_B(j,i, 1)**2

print(ans_W, ans_B)
  • DFS로 풀었다. 단, 내 병사(W)와 상대 병사(B)를 따로 탐색해주어야 하기에 따로 DFS와 visited 리스트를 만들어서 탐색했다.
  • 1차 탐색이 끝난 후(이어진 병사의 수 count) 제곱한 수를 ans에 더한다.

 


# 회고

 

그래프는 다른 탐색 이론들이랑 더 많이 사용되는 것 같다.

DP, 백트래킹, DFS/BFS 등을 더 연습하고 오자.

'코딩테스트' 카테고리의 다른 글

[Python] 조합 문제 풀이 정리(1)  (0) 2024.05.29
[Python] 조합  (0) 2024.05.25
[Python] 그래프 문제 풀이 정리(1)  (0) 2024.05.19
[Python] 트리  (0) 2024.05.17
[Python] 그래프(2)  (0) 2024.05.12