[파이썬] 백준 1944 : 복제 로봇 (골드1)

[파이썬] 백준 1944 : 복제 로봇 (골드1)

1944번: 복제 로봇

---

문제

---

풀이

0. 방향성 생각

• 읽어보면 모든 노드를 방문하는데 최소 비용을 쓴다.
• 각 노드에서 복제가 일어나서 각 노드별로 최단거리를 먼저 구해줘야함. 여기서 S=K라고 생각해도 좋다. 역할동일
• 이동 시 간선 비용이 같으니 BFS를 돌려서 각 노드별 거리를 구해준다.
• 모든 노드를 방문하면서, 간선 길이의 합이 최소가 되는 -> MST 문제임을 확인할 수 있다.

1. 입력

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

n,m = map(int,input().split())
arr = [list(input()) for _ in range(n)]
dire = [(1,0),(0,1),(-1,0),(0,-1)]

inf,count = 1e9,0
info = {}
for i in range(n):
    for j in range(n):
        if arr[i][j] == 'K' or arr[i][j] == 'S':
            info[(j,i)] = count
            count += 1

• 각 노드에 넘버링을 해주고 시작한다. (간선 정보에 인덱스 붙일거임)

2. BFS 함수 정의 + 간선 저장

def bfs(loc:tuple,node:int) -> list:
    x,y = loc
    visit = [[False]*n for _ in range(n)]
    visit[y][x] = True
    costs = [inf]*(m+1)
    costs[node] = 0
    q = deque([(x,y,0)])
    while q:    
        x,y,t = q.popleft()
        for dx,dy in dire:
            nx,ny,nt = x+dx,y+dy,t+1
            if 0<=nx<n and 0<=ny<n:
                na = arr[ny][nx]
                if na != '1' and not visit[ny][nx]:
                    q.append((nx,ny,nt))
                    visit[ny][nx] = True
                    if (na == 'K' or na == 'S') and nt < costs[info[(nx,ny)]]:
                        costs[info[(nx,ny)]] = nt 
    return costs
    
temp = []
for loc,num in info.items():
    temp.append(bfs(loc,num))
    
edges = []
for i in range(m+1):
    for j in range(m+1):
        if i<j and temp[i][j] != inf:
            edges.append((j,i,temp[i][j]))
edges.sort(key=lambda x:x[2],reverse=True)

• 좌표랑 노드 번호 받아서 다른 노드로 가는 거리를 리스트로 반환한다.
• 1만 아니면 이동할 수 있으니 큐에 추가하고 방문한다.
• 로봇, 키 노드에 도달한 경우 최소값으로 갱신해주기.
• temp에 간선을 추가해주고 edges에 이동가능한 간선만 뽑아준다 (inf가 아닌 경우)

3. find, union 정의 및 노드 연결

parent = list(range(m+1))
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

def union(x,y):
    parent[max(x,y)] = min(x,y)
    
answer,connect = 0,0
while edges:
    a,b,cost = edges.pop()
    pa,pb = find(a),find(b)
    if pa != pb:
        union(pa,pb)
        answer += cost
        connect += 1

if connect == m:
    print(answer)
else:
    print(-1)

• edges에서 부모가 다른 경우 union 진행.
• 로봇 1개 + 키 m개 = 총 m+1개 노드. 이 노드를 MST로 연결하려면 m개 간선이 필요
• 간선이 총 m개 이어졌으면 정답 출력. 아닌 경우 -1

 

---

전체코드

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

n,m = map(int,input().split())
arr = [list(input()) for _ in range(n)]
dire = [(1,0),(0,1),(-1,0),(0,-1)]

inf,count = 1e9,0
info = {}
for i in range(n):
    for j in range(n):
        if arr[i][j] == 'K' or arr[i][j] == 'S':
            info[(j,i)] = count
            count += 1

def bfs(loc:tuple,node:int) -> list:
    x,y = loc
    visit = [[False]*n for _ in range(n)]
    visit[y][x] = True
    costs = [inf]*(m+1)
    costs[node] = 0
    q = deque([(x,y,0)])
    while q:    
        x,y,t = q.popleft()
        for dx,dy in dire:
            nx,ny,nt = x+dx,y+dy,t+1
            if 0<=nx<n and 0<=ny<n:
                na = arr[ny][nx]
                if na != '1' and not visit[ny][nx]:
                    q.append((nx,ny,nt))
                    visit[ny][nx] = True
                    if (na == 'K' or na == 'S') and nt < costs[info[(nx,ny)]]:
                        costs[info[(nx,ny)]] = nt 
    return costs

temp = []
for loc,num in info.items():
    temp.append(bfs(loc,num))
    
edges = []
for i in range(m+1):
    for j in range(m+1):
        if i<j and temp[i][j] != inf:
            edges.append((j,i,temp[i][j]))
edges.sort(key=lambda x:x[2],reverse=True)

parent = list(range(m+1))
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])
    return parent[x]

def union(x,y):
    parent[max(x,y)] = min(x,y)
    
answer,connect = 0,0
while edges:
    a,b,cost = edges.pop()
    pa,pb = find(a),find(b)
    if pa != pb:
        union(pa,pb)
        answer += cost
        connect += 1

if connect == m:
    print(answer)
else:
    print(-1)

 

코멘트

프로그래머스 지형이동같은 느낌. graph 해석 후에 MST로

격자 내에서 이동 시 간선값이 다르거나 최적화가 들어갔으면 까다로웠을듯