bitsetの使い方

概要

bit演算子を勉強しようしよう思っていたのですが、以下の問題にぶち当たってようやくやる気になりました。bitsetの使い方がわかればbit演算子の威力がご理解いただけると思う。自分は思い知った。。。

問題

https://atcoder.jp/contests/abc147/tasks/abc147_e

bitsetの有用性

この問題は回答方針が決まってもTLEとなってしまう。最初はデータの保存をsetで管理していたのが問題だった。ACの人のコードを見てみると、bit演算を使って何やらかっこよく記述していた。尊敬の眼差しでした。自分なりの理解をした結果、bitsetでできることは、同じ値を足す、引くといった作業を繰り返す処理を一発でできるということ。ん、どういうこと？

bitset

その前に、bitsetについて簡単に解説する。bitsetは右からi番目（0からスタート）にその数値があるかないかを表す。百間は一見にしかずということで以下の例を見て欲しい。

数値	bitset
0	0
1	10
2	100
3	1000
[1, 3]	1010
[0, 2, 3]	1101

同じ値を足す、引く

このbitsetを使えば、同じ値を足す、引くという操作が一瞬で行える。例えば、以下のようなリストの各値に2を足したいとする。

a = [1, 2, 1, 4]

pythonで普通に書くと以下のようになる。リストの要素が大きくなるとかなり時間がかかってしまう。

for i in range(4):
    a[i] += 2

bitsetを使うと、この処理を一瞬でできる。

b = b << 2 #2を足す
b = b >> 2 #2を引く

コード

def run():
  ofs = 1000
  H, W = map(int, input().split())
  dp = [[0 for w in range(W)] for h in range(H)]
  a = [list(map(int, input().split())) for i in range(H)]
  b = [list(map(int, input().split())) for i in range(H)]
  c = [[abs(a[i][j] - b[i][j]) for j in range(W)] for i in range(H)]
  dp[0][0] = (1 << ofs + c[0][0]) | (1 << ofs - c[0][0])
  for h in range(1, H):
    dp[h][0] = (dp[h-1][0] << c[h][0]) | (dp[h-1][0] >> c[h][0])
  for w in range(1, W):
    dp[0][w] = (dp[0][w-1] << c[0][w]) | (dp[0][w-1] >> c[0][w])
  for h in range(1, H):
    for w in range(1, W):
      v = c[h][w]
      dp[h][w] = (dp[h-1][w] << v) | (dp[h-1][w] >> v)
      dp[h][w] |= (dp[h][w-1] << v) | (dp[h][w-1] >> v)
  s = list(bin(dp[-1][-1] >> ofs))
  for i in range(len(s)):
    if s[-1-i] == '1':
      print(i)
      break
 
if __name__ == '__main__':
  run()

参考文献

Python ビット演算 超入門
ABC147 E Balanced Pathでbitset高速化を完全に理解する
Balanced Path [AtCoder Beginner Contest 147 E]