貪欲法とDPで 1029. Two City Scheduling

問題はこちら。

Just a moment...

2N人をA,Bの2つのグループ分け、均等に振り分けるために必要なコストを最小化する問題。

Aに振り分けるコスト、Bに振り分けるコストが与えられる。

貪欲法

Aに行きやすい人はAに行かせ、Bに行きやすい人はBに行かせれば良い。

Aに行くコストとBに行くコストの差を取ってみて、差が大きければAに行くコストが大きいのでBに行かせた方が良いと考える。

コストの差の順に並び替え、AとBに振り分けていけばOK。

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        costs.sort(key=lambda x: x[0] - x[1])

        ans = 0
        n = len(costs)
        for i in range(n // 2):
            ans += costs[i][0] + costs[i + n // 2][1]

        return ans

計算量はソートがO(NlogN）、ループがO(N)となり、O(NlogN）となる。

メモリは答えを格納する変数ansだけなのでO(1)。

これが想定解。

DP

i-1人振り分けた時のコストが分かれば、i人目をAかBに振り分ける時のコストが分かるため、動的計画法（DP)が使えそうな感じがしたので考えてみた。

合計何人振り分けたか、Aに何人か、Bに何人か、という情報が必要となるので、3次元配列をDPで更新していく。

Aにa人、Bにb人、合計i人振り分けた時の状態をdp[i][a][b]として解いていく。

まずは配列を用意する。

コストを最小化する問題なので、初期値は無限大。

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        n = len(costs) // 2

        inf = float('inf')
        dp = [[[inf] * (n + 1) for _ in range(n + 1)] for _ in range(2 * n + 1)]

次に初期値の設定。

合計0人の所は誰もいないので0になる。

        for i in range(n + 1):
            for j in range(n + 1):
                dp[0][i][j] = 0

Aにだけ振り分ける、Bにだけ振り分ける時も初期値として計算可能。

累積和を求める要領でDPを更新。

        for i in range(1, 2 * n):
            for a in range(1, n + 1):
                dp[i][a][0] = dp[i - 1][a - 1][0] + costs[i - 1][0]
            for b in range(1, n + 1):
                dp[i][0][b] = dp[i - 1][0][b - 1] + costs[i - 1][1]

これで初期化が出来た。

メインルーチンを考える。

Aにa人、Bにb人、合計i人振り分けた時の状態をdp[i][a][b]とすると、

合計i-1人、Aにa-1人、Bにb人いる状態からAに1人追加する、または

合計i-1人、Aにa人、Bにb-1人いる状態からBに1人追加する、と遷移できる。

        for i in range(1, 2 * n + 1):
            for a in range(1, n + 1):
                for b in range(1, n + 1):
                    dp[i][a][b] = min(dp[i - 1][a - 1][b] + costs[i - 1][0],
                                      dp[i - 1][a][b - 1] + costs[i - 1][1])

求める答えは2N人を均等に振り分けた状態。

        ans = dp[2 * n][n][n]

        return ans

全部つなげると以下の通り。

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        n = len(costs) // 2

        inf = float('inf')
        dp = [[[inf] * (n + 1) for _ in range(n + 1)] for _ in range(2 * n + 1)]

        for i in range(n + 1):
            for j in range(n + 1):
                dp[0][i][j] = 0

        for i in range(1, 2 * n):
            for a in range(1, n + 1):
                dp[i][a][0] = dp[i - 1][a - 1][0] + costs[i - 1][0]
            for b in range(1, n + 1):
                dp[i][0][b] = dp[i - 1][0][b - 1] + costs[i - 1][1]

        for i in range(1, 2 * n + 1):
            for a in range(1, n + 1):
                for b in range(1, n + 1):
                    dp[i][a][b] = min(dp[i - 1][a - 1][b] + costs[i - 1][0],
                                      dp[i - 1][a][b - 1] + costs[i - 1][1])

        ans = dp[2 * n][n][n]

        return ans

計算量は前処理にO(N^2)が2回、DPの遷移がO(N^3）でO(N^3)となる。

メモリはDP配列の分必要なのでO（N^3)。

貪欲法に比べると非効率的な解法となった。