[파이썬] 백준 20303: 할로윈의 양아치 (골드3)

[파이썬] 백준 20303: 할로윈의 양아치 (골드3)

20303번: 할로윈의 양아치

---

문제

---

풀이

0. 방향성 생각

• BFS로 군집 크기, 사탕 수 계산
• DP로 최대값 구하기

1. 입력

from collections import deque
import sys
input = lambda : sys.stdin.readline().rstrip()

n,r,limit = map(int,input().split())
candy_info = dict(zip(range(1,n+1),list(map(int,input().split()))))

relations = [[] for _ in range(n+1)]
for _ in range(r):
    a,b = map(int,input().split())
    relations[a].append(b)
    relations[b].append(a)

 

2. BFS 함수 정의 + 탐색

visit = [False]*(n+1)
def bfs(x):
    candy,people = candy_info[x],1
    q = deque([x])
    visit[x] = True
    while q:
        x = q.popleft()
        for nx in relations[x]:
            if not visit[nx]:
                visit[nx] = True
                q.append(nx)
                people += 1
                candy += candy_info[nx]
    return (people,candy)

• 탐색하지 않은 node에서 BFS 진행
• 군집의 사람 수, 사탕 수 리턴하기

 

group_info = []
for i in range(1,n+1):
    if not visit[i]:
        group_info.append(bfs(i))

• group info에 리턴한 값들을 저장한다.

 

3. DP

dp = [[0]*limit for _ in range(len(group_info)+1)]
for i in range(len(group_info)):
    people,candy = group_info[i]
    for j in range(1,limit):
        if j>=people: dp[i+1][j] = max(dp[i][j],dp[i][j-people]+candy)
        else: dp[i+1][j] = dp[i][j]
print(max(dp[-1]))

• 그룹이 x개 일 때, x+1 크기의 dp를 만든다
• 아이들 무리를 1개 털 때 마다, 이전 정보들을 활용해 DP 테이블을 업데이트한다.
• i+1번 째 가방을 훔칠 때, 터는 무리의 사람 수 people이  j(누적 인원 수)보다 크면 누적 값을 갱신하게된다.
• dp[i][j-people]에서 사탕을 훔친 것과, 현재 들고있는 사탕 수 중 큰 값으로 갱신.
• 그렇지 않은 경우에는 이전 값을 그대로 올려보낸다.

---

전체코드

from collections import deque
import sys
input = lambda : sys.stdin.readline().rstrip()

n,r,limit = map(int,input().split())
candy_info = dict(zip(range(1,n+1),list(map(int,input().split()))))

relations = [[] for _ in range(n+1)]
for _ in range(r):
    a,b = map(int,input().split())
    relations[a].append(b)
    relations[b].append(a)

visit = [False]*(n+1)
def bfs(x):
    candy,people = candy_info[x],1
    q = deque([x])
    visit[x] = True
    while q:
        x = q.popleft()
        for nx in relations[x]:
            if not visit[nx]:
                visit[nx] = True
                q.append(nx)
                people += 1
                candy += candy_info[nx]
    return (people,candy)

group_info = []
for i in range(1,n+1):
    if not visit[i]:
        group_info.append(bfs(i))

dp = [[0]*limit for _ in range(len(group_info)+1)]
for i in range(len(group_info)):
    people,candy = group_info[i]
    for j in range(1,limit):
        if j>=people: dp[i+1][j] = max(dp[i][j],dp[i][j-people]+candy)
        else: dp[i+1][j] = dp[i][j]
print(max(dp[-1]))

 

코멘트

굳이 부모 갱신하면서 유니온 파인드 쓸 필요는 없는듯