[리트코드 3970] 연속 같은 문자 제한 최단 경로 — 연속 길이를 상태에 추가한 다익스트라

원제: Shortest Path With At Most K Consecutive Identical Characters

문제 요약

0번 노드에서 n - 1번 노드까지 가는 경로 중, 경로를 따라 방문한 노드의 labels를 이어 붙였을 때 같은 문자가 연속으로 k개를 초과하지 않는 경로를 찾습니다. 그런 경로가 있으면 간선 가중치 합의 최솟값을 반환하고, 없으면 -1을 반환합니다.

조건과 입출력

노드는 0부터 n - 1까지 번호가 붙은 방향 가중치 그래프입니다. edges[i] = [u_i, v_i, w_i]는 u_i에서 v_i로 가는 가중치 w_i의 간선이며, u_i != v_i입니다. labels의 길이는 n이고 소문자 영어로 이루어집니다. 1 <= n == labels.length <= 5 * 10^4, 0 <= edges.length <= 5 * 10^4, 1 <= w_i <= 10^4, 1 <= k <= 50입니다.

예제와 문제 자료

예제 1: n = 3, edges = [[0,1,1],[1,2,1],[0,2,3]], labels = "aab", k = 1이면 반환값은 3입니다. 0→1→2 경로의 라벨은 aab라서 연속 제한을 넘으므로, 0→2 경로의 비용 3을 선택합니다. 예제 2: 같은 그래프에서 k = 2이면 반환값은 2입니다. 라벨 aab가 허용되어 0→1→2 경로를 사용할 수 있습니다. 예제 3: n = 3, edges = [[0,1,1],[1,2,1]], labels = "aaa", k = 2이면 반환값은 -1입니다. 유일한 경로에서 같은 문자가 3개 연속됩니다.

풀이 아이디어

현재 노드만으로는 다음 이동이 유효한지 알 수 없습니다. 지금까지 이어진 같은 문자 개수도 필요하므로 다익스트라의 상태를 (node, cnt)로 둡니다. dist[node][cnt]는 이 상태에 도달하는 최소 비용입니다. 간선을 따라 같은 라벨로 이동하면 연속 개수가 늘고, 다른 라벨이면 1로 초기화합니다. 개수가 k를 넘는 이동은 이후에도 이미 성립한 연속 제한을 되돌릴 수 없으므로 버려도 됩니다. 가중치가 양수이므로 힙에서 유효한 최소 비용 상태를 꺼낼 때 목적지라면 최소 비용으로 도달한 것입니다.

코드의 핵심

코드는 방향 간선을 graph[u]에 (v, w)로 저장합니다. 시작 라벨도 연속 구간에 포함하므로 dist[0][1] = 0에서 시작하고 힙에도 (0, 0, 1)을 넣습니다.

간선 이동 시 labels[node] == labels[nxt]이면 next_cnt = cnt + 1, 다르면 next_cnt = 1입니다. next_cnt > k이면 그 이동은 건너뜁니다. 허용되는 경우에만 next_cost = cost + weight를 계산하고, next_cost < dist[nxt][next_cnt]일 때 상태를 갱신해 힙에 넣습니다.

예제 1에서 시작 상태는 (0, 0, 1)입니다. 0→1은 두 노드 라벨이 모두 a라 연속 개수가 2가 되어 k=1을 넘으므로 제외됩니다. 0→2는 라벨이 달라 연속 개수가 1이고 비용 3으로 도달합니다. 힙에서 목적지 상태가 꺼내져 3을 반환합니다.

힙에서 꺼낸 cost가 dist[node][cnt]보다 크면 이미 더 나은 경로로 갱신된 낡은 항목이므로 건너뜁니다. 목적지에 도달하는 상태가 없으면 반복이 끝난 뒤 -1을 반환합니다.

제출 코드

import heapq

class Solution:
    def shortestPath(self, n: int, edges: List[List[int]], labels: str, k: int) -> int:
        graph = [[] for i in range(n)]

        for u, v, w in edges:
            graph[u].append((v, w))

        dist = [[float('inf')] * (k + 1) for i in range(n)]
        dist[0][1] = 0

        heap = [(0, 0, 1)]

        while heap:
            cost, node, cnt = heapq.heappop(heap)

            if cost > dist[node][cnt]:
                continue

            if node == n - 1:
                return cost

            for nxt, weight in graph[node]:
                if labels[node] == labels[nxt]:
                    next_cnt = cnt + 1
                else:
                    next_cnt = 1

                if next_cnt > k:
                    continue

                next_cost = cost + weight

                if next_cost < dist[nxt][next_cnt]:
                    dist[nxt][next_cnt] = next_cost
                    heapq.heappush(heap, (next_cost, nxt, next_cnt))

        return -1

복잡도

상태 수는 최대 n(k + 1)입니다. 각 상태에서 나가는 간선을 확인하므로 전체 시간은 힙 연산을 포함해 O((n + |edges|)k log(|edges|k))로 볼 수 있습니다. dist와 그래프의 보조 공간은 O(nk + |edges|)이며, 힙은 지연된 항목을 포함해 최악의 경우 O(|edges|k) 공간을 사용할 수 있습니다.

복습 포인트

시작 노드의 라벨도 연속 개수에 포함하므로 초기 cnt는 1입니다. n = 1이면 시작 상태를 꺼낸 뒤 목적지 조건에 따라 0을 반환합니다. dist의 두 번째 인덱스는 0부터 k까지 만들지만, 이 코드에서는 유효한 상태의 연속 개수가 1 이상입니다.

출처

문제 원문