[Python] bit全探索 ZONE2021

ZONE2021C

チームメンバーの選び方を全探索すると、選び方が最大で $4.5×10^9$ 個になるため TLE してしまいます。

最大値の最小値のような形の問題では、答えを二分探索し、「答えが x 以上になるか？」という判定問題に持ち込むと簡単になることがあります。
この判定問題について考えましょう。

各能力値は、x 以上かどうかにしか興味がないので、x 以上なら1 、x 未満なら 0 とデータ圧縮することができます。すると、能力値の組としてあり得るものが $2^5=32$ 種類しかなくなるので、重複を取り除いた後 3 人の選び方を全探索することができます。

計算量は、能力値の種類を M 、採用する人数を K 、二分探索の上限を X として、$O((NM+2^{MK})logX)$ です。

個人的には、解法の方向性までは検討が付いたのだが、実装は歯が立たなかった。C問題としてはちょっと難しすぎる気がするが・・・。

サンプルコード

N = int(input())
A = [tuple(map(int, input().split())) for i in range(N)]

# x以上check
def check(x):
    s = set()
    for a in A:
        # x以上なら1を記録
        s.add(sum(1 << i for i in range(5) if a[i] >= x))
    for x in s:
        for y in s:
            for z in s:
                # 11111と等しいか、つまり各能力値がx以上か
                if x | y | z == 31:
                    return True
    return False

# 二分探索
ok = 0
ng = 10**9 + 1
while ng - ok > 1:
    cen = (ok + ng) // 2
    if check(cen):
        ok = cen
    else:
        ng = cen
print(ok)