[LeetCode] 3243. Shortest Distance After Road Addition Queries I - Python

 

 

난이도: medium

 

문제 설명

---

You are given an integer n and a 2D integer array queries.

 

There are n cities numbered from 0 to n - 1. Initially, there is a unidirectional road from city i to city i + 1 for all 0 <= i < n - 1.

 

queries[i] = [ui, vi] represents the addition of a new unidirectional road from city ui to city vi. After each query, you need to find the length of the shortest path from city 0 to city n - 1.

 

Return an array answer where for each i in the range [0, queries.length - 1], answer[i] is the length of the shortest path from city 0 to city n - 1 after processing the first i + 1 queries.

 

문제 예제

---

Example 1:

Input: n = 5, queries = [[2,4],[0,2],[0,4]]

Output: [3,2,1]

Explanation:

 

After the addition of the road from 2 to 4, the length of the shortest path from 0 to 4 is 3.

 

After the addition of the road from 0 to 2, the length of the shortest path from 0 to 4 is 2.

 

After the addition of the road from 0 to 4, the length of the shortest path from 0 to 4 is 1.

 

Example 2:

Input: n = 4, queries = [[0,3],[0,2]]

Output: [1,1]

Explanation:

 

After the addition of the road from 0 to 3, the length of the shortest path from 0 to 3 is 1.

 

After the addition of the road from 0 to 2, the length of the shortest path remains 1.

 

제한 사항

---

• 3 <= n <= 500
• 1 <= queries.length <= 500
• queries[i].length == 2
• 0 <= queries[i][0] < queries[i][1] < n
• 1 < queries[i][1] - queries[i][0]
• There are no repeated roads among the queries.

 

✏️ Solution(솔루션)

---

def dijkstra(graph, n, start):
    distances = [float(inf) for _ in range(n)]
    distances[start] = 0
    pq = [[0, start]]

    while pq:
        current = heapq.heappop(pq)
        if current[0] > distances[current[1]]:
            continue
        for to, distance in graph[current[1]]:
            now_distance = current[0] + distance
            if now_distance < distances[to]:
                distances[to] = now_distance
                heapq.heappush(pq, [now_distance, to])
    
    return distances[-1]

class Solution:
    def shortestDistanceAfterQueries(self, n: int, queries: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]
        answer = []

        for i in range(1, n):
            graph[i-1].append([i,1])
        
        for u, v in queries:
            graph[u].append([v,1])
            answer.append(dijkstra(graph, n, 0))
        
        return answer

 

나는 다익스트라 알고리즘을 통해 이 문제를 해결했다.

 

먼저 각 i번째 vertex가 i+1 vertex와 이어지는 그래프를 생성하기 위해 다음과 같이 코드를 작성했다.

 

for i in range(1, n):
	graph[i-1].append([i,1])

 

그 후 다음과 같이 queries에 있는 query 별로 graph에 edge를 추가해주었고, 추가한 다음에 다익스트라 알고리즘을 통해 vertex 0에서 vertex n-1 까지의 최단 경로를 구해 answer에 저장해주었다.

 

for u, v in queries:
	graph[u].append([v,1])
	answer.append(dijkstra(graph, n, 0))

 

이렇게 정답을 맞출 수 있었다.

 

다익스트라 알고리즘을 알고있으면 확실히 이런 최단경로를 찾는 문제를 쉽게 해결할 수 있는 것 같다.

 

사실 생각해보니 0에서 n-1까지의 최단경로를 구하면 되어서 dijkstra 알고리즘을 사용하여 0 에서 모든 vertex까지의 최단 경로를 구하지 않아도 될 것 같긴 하다.

 

---

[LeetCode] 3243. Shortest Distance After Road Addition Queries I