bitDPの初歩の次（2次元dp配列）【競技プログラミング】

はじめに

• bitDPについて初歩的(dp配列が1次元)なことを書いたけど、当然それだけでは身につかないから、dp配列が2次元になるケースを解いてみる。

問題例①

• 今回、問題例は生成AIに作って貰いました。
• 「bitDPの学習をしたいから、swiftでコード例と解説書いて」とお願いしたら、思ったのと違ったから、「さっきの例として、巡回セールスマン問題を使って」と付け加えたら、良い感じになりました。ただ、そのままだとエラーになったりしたので、若干修正しています。
• よくあるbitDPの解説は、c++かpythonで解説されているから、生成AIにswiftで作ってもらえるのは助かるね。
• ちなみに、巡回セールスマン問題とは、n個の都市と、それぞれの都市を結ぶルートの距離が与えられ、全ての都市を回って最初の都市に戻る最短ルートの距離を求めることであり、この問題では都市の数n=4、都市間の距離を二次元配列distとしている。

// import Foundation -- 不要のため削除

let n = 4 // 都市の数
let inf = Int.max / 2 // 2で割っておかないとエラーになるため「/ 2」を追記

// 距離行列
let dist = [
    [0, 10, 15, 20],
    [10, 0, 35, 25],
    [15, 35, 0, 30],
    [20, 25, 30, 0]
]

// dp配列
var dp = Array(repeating: Array(repeating: inf, count: n), count: 1 << n)

// 初期化
dp[1][0] = 0

// DPテーブルの更新
for mask in 1..<1 << n {
    for u in 0..<n {
        if mask & (1 << u) != 0 {
            for v in 0..<n {
                if mask & (1 << v) == 0 {
                    dp[mask | (1 << v)][v] = min(dp[mask | (1 << v)][v], dp[mask][u] + dist[u][v])
                }
            }
        }
    }
}

// 最短経路を計算
var result = inf
for u in 1..<n {
    result = min(result, dp[(1 << n) - 1][u] + dist[u][0])
}

print("最短経路の長さは \(result) です。")

• 上記のdp[mask][u]は、[mask:bit表現された訪問済みの都市][u:最後に訪問した都市]の最短距離が格納されています。
• 初期化では、訪問済み都市が「都市0」のみの1つだけという前提で、移動はしていないから、dp[1][0] = 0としています。
• DPテーブルの更新では、forループが3重になってる。
  • maskは訪問済みの都市を表すbit表現。
  • uは最後に訪問した都市。訪問済みの都市に含まれていることをチェック。
  • vはuの次に訪問する都市。訪問済みの都市に含まれていないことをチェック。
• 遷移式の肝の部分は、dp[mask | (1 << v)][v] ←← dp[mask][u] + dist[u][v]となっているけど、これは分かるよね。
  • 遷移後のbit表現について上記は、mask | (1 << v)って書いているけど、これはmask + 1<<vと同じです。まあ、それは分かるか。
• 配るdpで解いているけど、1次元dpを使って解いた前回と同じだね。
• dp配列を埋めた後、回答する際はスタート都市に戻る必要があるので、最後に訪問した都市をuとして、「dp[全て訪問済み][u]+dist[u][0]」が最も短くなるものを選ばないと駄目です。
• うむ、この問題も、配るdpとbit全探索を予めやっておけば難しい話じゃないね。

問題例②

• bitDPに取り組む原因となった教育的DPのo問題を解く。
• N人ずつ男女がいて、男iと女jの相性はAij(0または1)で表される。このとき、全ての男女のペアの相性を1とするマッチングパターンはいくつあるか？ただし、(mod 1000000007)で回答することとする。
• 入力例

3 // 3人ずつの男女
0 1 1 //男0と女0の相性は0、男0と女1の相性は1、男0と女2の相性は1
1 0 1
1 1 1

• 二次元dp[i][bf]は、男性陣を0..<iまで選んだとき(このとき、男性はi人選ばれている)、女性陣bfとペアを組ませるマッチングパターンとします。ここで、1ずつ遷移(増加)していくのは男性陣メンバーで、男性陣メンバーを増やすたびに、女性陣bf(bit表現集合)をぐるぐる回すイメージです。問題例①では、bit表現集合が遷移する主体で、追加する都市をぐるぐる回していたから、逆パターンですね。
• 初期状態は、男女を誰も選んでいない状態dp[0][0] = 1です。パターン数が0じゃなくて1かよ！って思っちゃうけど、累乗計算で$10^0$が1になるようなイメージで考えるしかない。深く考えたら負け。

extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
    init<T>(_ cnt2:Int, _ cnt1:Int, _ v:T) where Element == [T] {
        self = [[T]](cnt2,[T](cnt1,v))
    }
}

let N = Int(readLine()!)!

//2次元配列
var A:[[Int]] = []
for _ in 0..<N {
    let vs = readLine()!.split(separator:" ").map{Int($0)!}
    A.append(vs)
}

var dp = [[Int]](N+1,1<<N,0)
dp[0][0] = 1

// 遷移（配るdp）
for i in 0..<N {
    for bf in 0..<(1<<N) {
        if i == bf.nonzeroBitCount {
            for j in 0..<N {
                if bf & 1<<j == 0 && A[i][j] == 1 {
                    dp[i+1][bf + 1<<j] += dp[i][bf] // 遷移式
                }
            }
        }
    }
}

print(dp[N][1<<N - 1])

• 上記コードで遷移式はdp[i+1][bf + 1<<j] += dp[i][bf]だけど、これも配るdpに分類して良いのかな？一番内側のforループで変化するjが遷移式の左辺にいるから、僕の理解では配るdpです。
• 遷移は三重forループだけど、途中でifチェックを行っています。
  • i == bf.nonzeroBitCount
    • 男性がi人選ばれているので、女性陣集合を表すbfで立っているビットの数がiの時だけチェックを通す。
  • bf & 1<<j == 0 && A[i][j] == 1
    • 女性jが、追加前集合bfにまだ含まれておらず、また、男性iと女性jの相性が1の時のみ、チェックを通す。
• 上記コードを提出すると...おっと、危ない。mod 1000000007を忘れてた。遷移式の次の行に
  dp[i+1][bf + 1<<j] %= 1000000007を追加してから提出すると、650ms程度でAC。

さいごに

• bitDPにも慣れてきた！そして、生成AIって、意外と優秀なのね。数年前と大違いでビックリした。
• サラリーマンの仕事のうち、単調作業部分は確実に楽になるね。今回の投稿では、bidDPの学習自体より、生成AIの優秀さへの驚きの方が心に刻まれちゃったよ。