Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 두 지점의 금을 줍는 최소 비용 구하기 – 다익스트라 알고리즘 완벽 정리

문제 개요

2차원 행렬(matrix)과 여러 값들이 주어집니다: row, col, erow0, ecol0, erow1, ecol1. 현재 우리의 위치는 matrix[row, col]이며, 이곳에서 출발해 matrix[erow0, ecol0]matrix[erow1, ecol1] 두 곳에 있는 금을 모두 줍고자 합니다.

이동은 상하좌우 네 방향으로 자유롭게 할 수 있지만, 셀(r, c)에 도착할 때마다 해당 셀의 비용인 matrix[r, c]를 지불해야 합니다. 단, 같은 셀을 여러 번 방문하더라도 비용은 한 번만 지불합니다. 목표는 두 지점의 금을 모두 줍기 위한 최소 비용을 구하는 것입니다.

예시로 이해하기

입력이 다음과 같다고 가정해 보겠습니다.

11111
110101010
1111010

조건: row = 0, col = 0, erow0 = 0, ecol0 = 3, erow1 = 2, ecol1 = 2

이 경우 출력은 8입니다. 현재 위치 (0, 0)에서 금이 있는 (0, 3)과 (2, 2)로 이동해야 하는데, 먼저 오른쪽으로 세 칸 이동해 (0, 3)에 도착한 뒤 되돌아와서, 비용이 1인 셀들을 따라 (2, 2)까지 내려가면 최적 경로가 됩니다.

풀이 접근 방법

이 문제는 각 시작점에서 모든 셀까지의 최단 거리를 구하는 다익스트라(Dijkstra) 알고리즘으로 해결할 수 있습니다. 핵심 아이디어는 세 지점(현재 위치, 첫 번째 금 위치, 두 번째 금 위치) 각각에서 다익스트라를 수행한 뒤, 세 경로가 만나는 최적의 중간 지점을 찾는 것입니다.

알고리즘 단계

  • is_valid(x, y) 함수를 정의합니다. 좌표 x, y가 행렬 범위 안에 있으면 true, 아니면 false를 반환합니다.

  • min_cost(sx, sy) 함수를 정의합니다. 시작점 (sx, sy)에서 다익스트라를 수행합니다.

    • heap := (matrix[sx, sy], sx, sy)를 담은 우선순위 큐로 초기화

    • dists := 주어진 행렬과 같은 크기의 배열, 무한대(inf)로 채움

    • dists[sx, sy] := matrix[sx, sy]

    • heap이 빌 때까지 반복:

      • (cost, x, y) := heap에서 최솟값 추출 후 제거

      • 네 방향 인접 셀 (nx, ny) 각각에 대해:

        • is_valid(nx, ny)이고 matrix[nx, ny] + cost < dists[nx, ny]이면:

          • edge := matrix[nx, ny]

          • dists[nx, ny] := edge + cost

          • (edge + cost, nx, ny)를 heap에 삽입

    • dists 반환

  • 메인 로직에서 다음을 수행합니다:

    • res := 무한대(inf)

    • a := min_cost(row, col), b := min_cost(erow0, ecol0), c := min_cost(erow1, ecol1)

    • 모든 셀 (i, j)에 대해 res := min(res, a[i, j] + b[i, j] + c[i, j] − 2 × matrix[i, j])

    • res 반환

왜 −2 × matrix[i, j]를 빼나요?

세 개의 최단 거리 배열(a, b, c)은 각각 자신의 시작점 비용을 포함하고 있으며, 만나는 지점 (i, j)의 비용은 세 배열 모두에 중복해서 포함됩니다. 따라서 중복 계산된 만남 지점의 비용 2번분을 빼주어야 실제 총비용이 정확하게 계산됩니다.

구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

import heapq
import math
class Solution:
   def solve(self, matrix, row, col, erow0, ecol0, erow1, ecol1):
      def is_valid(x, y):
         return x >= 0 and y >= 0 and x < len(matrix) and y < len(matrix[0])
      def min_cost(sx, sy):
         heap = [(matrix[sx][sy], sx, sy)]
         dists = [[math.inf] * len(matrix[0]) for _ in range(len(matrix))]
         dists[sx][sy] = matrix[sx][sy]
         while heap:
            cost, x, y = heapq.heappop(heap)
            for nx, ny in [(x, y - 1), (x + 1, y), (x - 1, y), (x, y + 1)]:
               if is_valid(nx, ny) and matrix[nx][ny] + cost < dists[nx][ny]:
                  edge = matrix[nx][ny]
                  dists[nx][ny] = edge + cost
                  heapq.heappush(heap, (edge + cost, nx, ny))
         return dists
      res = math.inf
      a, b, c = min_cost(row, col), min_cost(erow0, ecol0), min_cost(erow1, ecol1)
      for i in range(len(matrix)):
         for j in range(len(matrix[0])):
            res = min(res, a[i][j] + b[i][j] + c[i][j] - 2 * matrix[i][j])
      return res
ob = Solution()
matrix = [
   [1, 1, 1, 1, 1],
   [1, 10, 10, 10, 10],
   [1, 1, 1, 10, 10]
]
row = 0
col = 0
erow0 = 0
ecol0 = 3
erow1 = 2
ecol1 = 2
print(ob.solve(matrix, row, col, erow0, ecol0, erow1, ecol1))

입력

[ [1, 1, 1, 1, 1], [1, 10, 10, 10, 10], [1, 1, 1, 10, 10]
], 0, 0, 0, 3, 2, 2

출력

8

마무리

이 문제의 시간 복잡도는 다익스트라를 세 번 수행하므로 O(N×M×log(N×M))입니다(N, M은 행렬의 크기). 우선순위 큐(heapq)를 활용하면 그리드 형태의 가중치 최단 경로 문제를 효율적으로 해결할 수 있으며, 이 패턴은 여러 출발점 사이의 최적 만남 지점을 찾는 다양한 변형 문제에도 응용할 수 있습니다.