[파이썬] 1949. 우수 마을

coding test

[파이썬] 1949. 우수 마을

잔망루피 2021. 5. 20. 12:29

문제

N개의 마을로 이루어진 나라가 있다. 편의상 마을에는 1부터 N까지 번호가 붙어 있다고 하자. 이 나라는 트리(Tree) 구조로 이루어져 있다. 즉 마을과 마을 사이를 직접 잇는 N-1개의 길이 있으며, 각 길은 방향성이 없어서 A번 마을에서 B번 마을로 갈 수 있다면 B번 마을에서 A번 마을로 갈 수 있다. 또, 모든 마을은 연결되어 있다. 두 마을 사이에 직접 잇는 길이 있을 때, 두 마을이 인접해 있다고 한다.

이 나라의 주민들에게 성취감을 높여 주기 위해, 다음 세 가지 조건을 만족하면서 N개의 마을 중 몇 개의 마을을 '우수 마을'로 선정하려고 한다.

'우수 마을'로 선정된 마을 주민 수의 총 합을 최대로 해야 한다.
마을 사이의 충돌을 방지하기 위해서, 만일 두 마을이 인접해 있으면 두 마을을 모두 '우수 마을'로 선정할 수는 없다. 즉 '우수 마을'끼리는 서로 인접해 있을 수 없다.
선정되지 못한 마을에 경각심을 불러일으키기 위해서, '우수 마을'로 선정되지 못한 마을은 적어도 하나의 '우수 마을'과는 인접해 있어야 한다.

각 마을 주민 수와 마을 사이의 길에 대한 정보가 주어졌을 때, 주어진 조건을 만족하도록 '우수 마을'을 선정하는 프로그램을 작성하시오.

입력

첫째 줄에 정수 N이 주어진다. (1≤N≤10,000) 둘째 줄에는 마을 주민 수를 나타내는 N개의 자연수가 빈칸을 사이에 두고 주어진다. 1번 마을부터 N번 마을까지 순서대로 주어지며, 주민 수는 10,000 이하이다. 셋째 줄부터 N-1개 줄에 걸쳐서 인접한 두 마을의 번호가 빈칸을 사이에 두고 주어진다.

출력

첫째 줄에 '우수 마을'의 주민

수의 총 합을 출력한다.

예제 입력 1

1000 3000 4000 1000 2000 2000 7000

1 2

2 3

4 3

4 5

6 2

6 7

예제 출력 1

14000

🙋‍♀️ 나의 풀이

import sys, collections
sys.setrecursionlimit(10**5)

def dfs(cur):
    visited[cur]=1
    for u in g[cur] :
        if not visited[u] :
            dfs(u)
            dp[cur][1]+=dp[u][0]    # 현재 마을을 우수 마을로 선정
            dp[cur][0]+=max(dp[u][0], dp[u][1])     # 현재 마을을 우수 마을로 선정x

n=int(sys.stdin.readline().strip())
cost=[0]+[int(x) for x in sys.stdin.readline().split()]

visited=[0 for _ in range(n+1)]
dp=[[0, cost[i]] for i in range(n+1)]     # dp[i][0]=i 마을을 선정x, dp[i][1]=i마을을 선정
g=collections.defaultdict(list)    # 마을의 연결 정보를 담는 딕셔너리

for _ in range(n-1) :
    v, u=map(int, sys.stdin.readline().split())
    g[v].append(u)
    g[u].append(v)

dfs(1)
print(max(dp[1][1], dp[1][0]))

dfs와 dp 사용

풀다가 실패해서 다른 사람 풀이 참고함..

recursion 깊이를 10^4로 하면 런타임 에러(Recursion Error)가 뜬다.

N이 최대 10,000이기 때문에 넉넉하게 10^5가 좋다.

# 실패
N=int(input())    # 마을 갯수
people=list(map(int, input().split()))    # 마을 별 주민 수
dic={}                 # 딕셔너리
tree=[[] for _ in range(N+1)]
stack=list()
visited=[0]*(N+1)
ans=0

good_city=set()
un_city=set()


# key=마을번호, 값=주민 수
for i in range(1, N+1) :
    dic[i]=people[i-1]

all_num=set(dic.keys())

# 마을 연결하기
for i in range(N-1) :
    a, b=map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)

start=sorted(tree, reverse=True, key=lambda x : len(x))[0]  # 자식이 많은 순 정렬
_index=tree.index(start)        # 시작 번호
stack=[_index]

# 탐색
while stack :
    M = 0
    person=stack.pop()
    if not visited[person] :
        visited[person]=1
        for i in tree[person] :     # 자식들 중 가장 인구 많은 자식 찾기
            if not visited[i] :
                if M < dic[i] :
                    M=dic[i]     # 갱신
                    M_index=i           # 인덱스
        if dic[person] > M :     # 부모가 자식보다 인구가 많음
            ans+=dic[person]
            good_city.add(person)    # 우수 마을 번호 등록
            for t in tree[person]:
                un_city.add(t)    # 우수 마을과 연결된 노드들 x
        else :
            # 자식 스택에 넣기
            for j in tree[person]:
                stack.append(j)

# 덜 된 것
j=list(all_num-(good_city | un_city))
if j :
    for id in j :
        ans+=dic[id]
print(ans)

문제에 나온 예시는 통과하지만 제출하니 실패 ㅠ

스택을 이용해 dfs로 구현.

연결된 노드가 많은 노드부터 시작했다.

우수 마을에 노드를 넣으면서 우수 마을과 연결된 노드들은 un_city에 넣는다.

전체 all_num에서 우수 마을 집합 good_city와 우수 마을이 될 수 없는 집합 un_city에 속하지 않는 값은 우수 마을로 넣는다.

👩‍🦰 다른 사람 풀이

# https://cotak.tistory.com/29
import sys, collections
sys.setrecursionlimit(10**6)

def dfs(cur):
    visited[cur]=1
    for u in g[cur] :
        if not visited[u] :
            dfs(u)
            dp[cur][1]+=dp[u][0]    # 현재 마을을 우수 마을로 선정
            dp[cur][0]+=max(dp[u][0], dp[u][1])     # 현재 마을을 우수 마을로 선정x

n=int(sys.stdin.readline().strip())
cost=[0]+[int(x) for x in sys.stdin.readline().split()]

visited=[0 for _ in range(n+1)]
dp=[[0, cost[i]]*2 for i in range(n+1)]     # dp[i][0]=i 마을을 선정x, dp[i][1]=i마을을 선정
g=collections.defaultdict(list)

for _ in range(n-1) :
    v, u=map(int, sys.stdin.readline().split())
    g[v].append(u)
    g[u].append(v)

dfs(1)
print(max(dp[1][1], dp[1][0]))

dp와 재귀호출을 이용한 dfs 풀이

딕셔너리 g에 마을의 연결 정보를 양방향으로 넣는다.

dp=[[0, cost[i]]*2 for i in range(n+1)] 에서 왜 [0, cost[i]]*2를 하나 이해가 안 갔는데 *2 없는 게 맞음.

cur을 우수 마을로 선정하는 것은 u(cur과 인접한 마을)를 선정하지 않는 것이다.

cur을 우수 마을로 선정하지 않는 것은 u를 우수 마을로 선정하는 것과 하지 않는 것 2가지를 고려해야 한다.

cur에 연결된 u가 몇 개인지 모르니까 두 가지 경우 중 더 큰 쪽을 포함한다.

# https://conak-diary.tistory.com/107
def DFS(current):
    visited[current]=True
    s=[current]
    DP[current][0]=population[current]
    while s :
        current=s[-1]
        for next in connection[current] :
            if not visited[next] :
                s.append(next)
                visited[next]=True
                DP[next][0]=population[next]
                break
        else :	# break가 실행되지 않았으면
            s.pop()
            if s:
                DP[s[-1]][0] += DP[current][1]
                DP[s[-1]][1] += max(DP[current][0], DP[current][1])


N=int(input())
population=[0]+list(map(int, input().split()))
connection=[[] for _ in range(N+1)]

for i in range(N-1) :
    a, b=map(int, input().split())
    connection[a].append(b)
    connection[b].append(a)

DP=[[0, 0] for _ in range(N+1)]
visited=[False]*(N+1)
DFS(1)
print(max(DP[1][0], DP[1][1]))

스택을 이용한 dfs

문제 출처 💁‍♀️ 백준

'coding test' 카테고리의 다른 글

[파이썬] 1303. 전쟁 - 전투 (0)	2021.05.22
[파이썬] 1260. DFS와 BFS (0)	2021.05.21
[파이썬] 2293. 동전 1 (0)	2021.05.20
[파이썬] 3085. 사탕 게임 (0)	2021.05.20
[파이썬] 2309. 일곱 난쟁이 (0)	2021.05.19

현재글[파이썬] 1949. 우수 마을

나의 공부 기록

Roll with the punches 🚵‍♂️

나의 공부 기록