[파이썬] 백준 1414 : 불우이웃돕기 (골드3)

[파이썬] 백준 1414 : 불우이웃돕기 (골드3)

1414번: 불우이웃돕기

---

문제

 

---

풀이

0. 방향성 생각

• MST

1. 입력

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

def cvt(x):
    nx = ord(x)
    if 65<=nx<91: return nx-38
    elif 97<=nx<123: return nx-96
    
n = int(input())
arr = [list(map(cvt,input())) for _ in range(n)]

edges = []
for i in range(n):
    for j in range(n):
        if arr[i][j]:
            edges.append(((j,i),arr[i][j]))
edges.sort(key=lambda x:-x[1])

• 아스키코드 변환해주기. map int처럼 한 번에 해주기

2. 유니온파인드 함수 정의

parent = list(range(n))
def find(a):
    if parent[a] != a:
        parent[a] = find(parent[a])
    return parent[a]
 
def union(a,b):
    pa,pb = find(a),find(b)
    if pa != pb:
        parent[max(pa,pb)] = min(pa,pb)

• 작은애를 부모로 지정해주기.

3. 출력

answer = 0
while edges:
    nodes,cost = edges.pop()
    x,y = nodes
    answer += cost
    if find(x) != find(y):
        union(x,y)
        answer -= cost
        
for i in range(n):
    find(i)
    
if len(set(parent)) != 1:
    print(-1)
else:
    print(answer)

• 간선 정렬하고 일단 기부할 길이 answer에 올려놓는다.
• 그 간선을 사용해야하면 다시 answer에서 빼주기.
• 유니온 파인드가 끝난 후, 모두 연결되어있는지 확인을 하려면 모든 노드에서 find를 진행해줘서 부모를 업데이트 해준다.
• 만약 parent 종류가 하나가 아니면, 모두 연결하지 못한 케이스임. 하나이면 기부 가능한 길이 출력

 

---

전체코드

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

def cvt(x):
    nx = ord(x)
    if 65<=nx<91: return nx-38
    elif 97<=nx<123: return nx-96
    
n = int(input())
arr = [list(map(cvt,input())) for _ in range(n)]

edges = []
for i in range(n):
    for j in range(n):
        if arr[i][j]:
            edges.append(((j,i),arr[i][j]))
edges.sort(key=lambda x:-x[1])

parent = list(range(n))
def find(a):
    if parent[a] != a:
        parent[a] = find(parent[a])
    return parent[a]
 
def union(a,b):
    pa,pb = find(a),find(b)
    if pa != pb:
        parent[max(pa,pb)] = min(pa,pb)

answer = 0
while edges:
    nodes,cost = edges.pop()
    x,y = nodes
    answer += cost
    if find(x) != find(y):
        union(x,y)
        answer -= cost
        
for i in range(n):
    find(i)
    
if len(set(parent)) != 1:
    print(-1)
else:
    print(answer)

 

코멘트

유니온 파인드가 끝난 후에 parent가 업데이트 안되는거는 생각 안하고있었는데 좋은 생각거리인듯