bitDPとメモ化再帰の合わせ技【競技プログラミング】

はじめに

• 「教育的dpコンテスト」の問題にペアの相性を得点化して、最大化する問題があった。
• 問題Uだけど、これは、メモ化再帰とbitDPの2つのワザを使っているので、復習がてら解いてみようと思う。

問題Uの内容

• ウサギがN匹います。ウサギ同士には相性があり、相性は二次元配列A[i][j]で与えられます。
  • 相性の条件
    • A[i][i]=0
      • 本来的には同じウサギ同士の相性の値は不要だけど、0にすることでコード上で除外する手間を省けるようにしてるのかもね
    • A[i][j]=A[j][i]
      • 相性だから、どちらも同値としてるんだろうね。愛情ならi⇒jとj⇒iで違う値かもね
• ウサギはどこかのグループに属します。なお、グループは1匹でも形成可能。
  • 例えば、3匹のウサギ(1 2 3)のグループ分け方法は、以下の5とおり。「グループ分けについてはコレを参考にしてくれよ、ベル君！」(cv:ドヤ○ング)
    • {1 2 3}
    • {1} {2} {3}
    • {1} {2 3}
    • {2} {1 3}
    • {3} {1 2}
• グループに2匹以上のウサギがいる場合は、各ウサギ同士の相性が得点に加算されます。
• 与えられた二次元配列Aで可能な最大得点を答えなさい。
  • もし、Aの要素に負数がなければ、全てのウサギを一つのグループにすれば得点は最大化されるね。
  • また、Aの要素にどれだけ負数が含まれていても、要素一つのグループを作れるから、最大得点がマイナスになることはないね。

問題Uの解き方

• まず、再帰関数rec(bf)の書き方です！
  • bfはビットフラグの意味として使い、例えば、ウサギ4匹[0 1 2 3]のうち、[0]だけいる場合はbf=0b0001、[0 1]がいるときは、bf=0b0011、全部いるときはbf=0b1111になるよ！
  • dp[bf]はメモ化再帰におけるメモだね！初期値をinf(=-Int.max / 2)にして再帰関数rec(bf)の答えを上書きするよ！

func rec(_ bf:Int)->Int {
    if dp[bf] != inf {return dp[bf]}
    var ans = pts(bf) // ptsは、bfをグループとするときの点数
    for ss in fss(bf) { // fssはbfの部分集合
        ans = max(ans , rec(ss) + rec(ss^bf))
    }
    dp[bf] = ans // メモ化
    return ans
}

• 関数ptsとfssは別途説明するけど、上記コードの意味は分かるかな？
• 分からない人は、スライム君を切り分ける解法を参考にしてください。
• この問題は、bitDPとメモ化再帰が分かれば、大して難しくないね！

関数ptsとfss

• 関数ptsは説明するまでもないよね。こんな感じです。名前は、得点をpointsと考えて、短縮してpts。

func pts(_ bf:Int)->Int{
    if bf.nonzeroBitCount < 2 {return 0}
    var sum = 0
    for i in 0..<N {
        for j in i+1..<N {
            if (bf>>i & 1) & (bf>>j & 1) == 0 {continue}
            sum += A[i][j]
        }
    }
    return sum
}

• 関数fssは後で説明するね。名前は、部分集合族（family of subset）を短縮してfss。
  • 冪集合(PowerSet)を短縮してpsにしようかとも思ったけど、今回は、元となる集合bfとゼロ集合を含まないで返すから、psは適切じゃないと考えたよ。

func fss(_ bf:Int)->[Int]{
    var ans:[Int] = []
    if bf == 0 {return []}
    var ss = bf
    while ss>0 {
        ss -= 1
        ss = ss&bf // ここがキモ！
        if ss == 0 {break}        
        ans.append(ss)
        // 参考表示
        // print(String(bin:bf),String(bin:ss),ss)
    }
    return ans
}

• fssは、実際にbf(二進数表示)を食わせてみて、どんなss(二進数表示)を吐き出すかを見れば理解できると思うよ。上記のコメントアウトしているprint文を使うよ。
• ちなみに、swiftは、Stringで二進数表示出来るけど、ゼロ埋めで桁を揃える事の出来ないアホアホ言語なので、extensionでStringに機能追加してるよ。（16進数ならformatで桁揃えできるのに、2進数だと出来ないのが解せぬ。）

extension String{ // 二進数表示機能を強制追加！
    init(bin:Int,len:Int = 4){
        let str = Self(bin,radix:2)
        let zeroLen = len - str.count < 0 ? 0 : len - str.count
        let zero = Self(repeating:"0",count:zeroLen)
        self = zero + str
    }
}

• fssにbf=13(0b1101)を食わせたときのprint文の実行結果です。左側がbfで、真ん中がss、右端がssの10進数表示だよ。

1101 1100 12
1101 1001 9
1101 1000 8
1101 0101 5
1101 0100 4
1101 0001 1

• 上のprint文の結果とコードを照らし合わせれば、やってることは分かるよね！
• このbfの部分集合を出す関数は意外と重要だと思うよ！
  • 冪集合(bfの部分集合全て)にしたいときは、「bf」と「0」が含まれるようにコードを書き換えてね。

問題Ｕの解答コード全体

extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
}

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

let S = 1<<N - 1 // 2^NのbitDP用のbf（ビットフラグ）の初期状態
let inf = -Int.max / 2
var dp = [Int](1<<N,inf)

func rec(_ bf:Int)->Int {
    if dp[bf] != inf {return dp[bf]}
    var ans = pts(bf) // ptsは、bfをグループとするときの点数
    for ss in fss(bf) { // fssはbfの部分集合
        ans = max(ans , rec(ss) + rec(ss^bf))
    }
    dp[bf] = ans // メモ化
    return ans
}

func pts(_ bf:Int)->Int{
    if bf.nonzeroBitCount < 2 {return 0}
    var sum = 0
    for i in 0..<N {
        for j in i+1..<N {
            if (bf>>i & 1) & (bf>>j & 1) == 0 {continue}
            sum += A[i][j]
        }
    }
    return sum
}

func fss(_ bf:Int)->[Int]{
    var ans:[Int] = []
    if bf == 0 {return []}
    var ss = bf
    while ss>0 {
        ss -= 1
        ss = ss&bf // ここがキモ！
        if ss == 0 {break}        
        ans.append(ss)
    }
    return ans
}

rec(S)
print(dp[S])

• コレで提出したら、300ms程度でACでした。

さいごに

• メモ化再帰とbitDPについては、体に馴染ませないとコンテストで戦えないと分かってきたよ。
• 「教育的dpコンテスト」の残り問題も僅かになってきた。最後まで頑張りたい。