leetcode 787 - Cheapest Flights Within K Stops

https://leetcode.com/problems/cheapest-flights-within-k-stops/

풀이

최대 k번 경유할 수 있기 때문에 아래와 같이 테이블을 정의한다.

d[i][j] = 최대 i번 경유해서 j까지 왔을 때 최단 거리

정의한 대로 d[i][j]를 채우기 위해 점화식을 세워보면 j로 갈 수 있는 모든 k에 대해서 아래와 같은 식이 된다.

d[i][j] = min(d[i-1][k] + distance(k, j))

 

물론 최대 i번 경유라고 정의했기 때문에 처음에 d[i][j] = d[i-1][j]로 초기화를 해두어야 한다.

식을 보면 i번 경유할때 참조하는 데이터는 i-1번 경유한 경우밖에 없다. 따라서 d 배열을 2차원 테이블로 만들 필요 없이 1차원 테이블 2개로 재활용하면서 구하면 메모리를 아낄 수 있다.

 

코드

class Solution {
public:
    int findCheapestPrice(int n, vector<vector<int>>& flights, int src, int dst, int K) {
        vector<vector<pair<int, int>>> graph(n, vector<pair<int,int>>());
        for (const auto& flight : flights) {
            graph[flight[0]].push_back({ flight[1], flight[2] });
        }

        vector<int> cost(n, INT_MAX);
        cost[src] = 0;

        for(int i = 0; i <= K; i++) {
            vector<int> next(cost);
            for(int j = 0; j < n; j++) {
                if(cost[j] == INT_MAX) continue;
                for(int k = 0; k < graph[j].size(); k++) {
                    next[graph[j][k].first] = min(next[graph[j][k].first], cost[j] + graph[j][k].second);
                }
            };
            cost = next;
        }

        return cost[dst] == INT_MAX ? -1 : cost[dst];
    }
};