qumazakiのブログ

pythonで競プロをしています。

AtCoder 東京海上日動 プログラミングコンテスト2020 E - O(rand)

問題解説

 

atcoder.jp 公式解説がわかりにくかったので自分なりに書いてみます。

問題概要

A_1,A_2,...,A_Nが与えられます。
ここから重複なく1～K個選び、下記の条件を満たす方法は何通りあるか求めよ。
・選んだ数の論理積がS
・選んだ数の論理和がT

制約

1 \leq K \leq N \leq 50 \\ 0 \leq A_i, S, T < 2^{18} \\ A_i \neq Aj

方針

  1. S,Tに関する問題の単純化
  2. 数え上げ方針の検討
  3. 数え上げの実装

1. S,Tに関する問題の単純化

まずS,Tについて観察します。
各ビットごとに見たとき、Sのi桁目がS_i、Tのi桁目がT_iだったとします。
すると、下記の場合分けができます。
①(S_i,T_i) = (0,0)
→　選んだ数は全てi桁目が0でなくてはならない。

②(S_i,T_i) = (0,1)
→　選んだ数のi桁目に0,1両方が混ざっている必要がある。

③(S_i,T_i) = (1,0)
→　このような選び方は存在しない。
　　→　この組み合わせがあった場合は答えが0通りで確定

④(S_i,T_i) = (1,1)
→　選んだ数は全てi桁目が1でなくてはならない。

③はなかったものとして、①と④を考えると、
i桁目の数字に着目することで選ぶことのできない数字を取り除くことができます。

例えば入力例2を見ると、これは1桁目が④、2,3桁目が②、4桁目以降が①になっている例です。

S = 1 \\ T = 7 \\ A = \{3,4,9,1,5\}

よって、1桁目が1ではないA_2 = 4は使うことができません。
同様に、4桁目が0ではないA_3 = 9も使うことができません。
逆にいうと、それ以外の数字を使えば1,4桁目の条件はいつでも満たすことができます。

以上のことから、使える数字だけを選定し、②以外の桁を右に詰めてしまうことにしましょう。
先の入力例2は下記のように変わります。

S = 0 \\ T = 3 \\ A = \{1,,,0,2\}

このような前処理をすることで、問題を下記のように言い換えることができました。

S=0,T=2^ L-1,0 \leq A_i \leq 2^ Lとして、1～K個の数を選ぶ。
各桁ごとに選んだ数のビットが0か1かをカウントしたとき、全ての桁で0,1両方が1回以上出現しているような選び方は何通りか
※Lは②を満たす桁の数

2. 数え上げ方針の検討

ここからが難しいです。
結論から言うと、桁ごとが満たすべき条件に着目し、包除原理を使って解きます。

まず、「全ての桁で0,1両方が1回以上出現しているような選び方」ですが、 これは各桁ごとに独立な事象の重ね合わせです。 つまり、下記で表現できます。

「1桁目で0,1両方が1回以上出現しているような選び方」
\cup「2桁目で0,1両方が1回以上出現しているような選び方」
\cup ...
\cup「L桁目で0,1両方が1回以上出現しているような選び方」

また、「0,1両方が1回以上出現しているような選び方」は扱いづらいです。
「0だけ」「1だけ」「両方」の3状態を考えないといけないので状態量が3^ Lになってしまいます。

ここで 、余事象を考えると扱いやすくなります（後述）。
言語化すると「"0だけが出現している"or"1だけが出現している"ような選び方」になります。

すると求める答えは下記に言い換えられます。

「1桁目で"0だけが出現している"or"1だけが出現している"ような選び方」
\cap「2桁目で"0だけが出現している"or"1だけが出現している"ような選び方」
\cap ...
\cap「L桁目で"0だけが出現している"or"1だけが出現している"ような選び方」
の余事象

これで包除原理が使えるような形に持ち込むことができました。

i桁目が「"0だけが出現している"or"1だけが出現している"ような選び方」を満たしているとき、X_i = 1、満たしていてもいなくてもどちらでも良いとき X_i = 0とします。
X := \{ X_1,...X_L \} | (X_i =0,1)を一つ決めたときに、
そのような条件を満たす数の選び方をf(X)で定めます。
Xの2^ L通りの選び方全体をYとすると、
（包除原理については詳しい説明は省きますが、）答えは下記で求めることができます。

\Large \Sigma_{X \in Y} \{(-1)^{\Sigma X_i} * f(X) \}

あとは、f(X)が求められれば良いです。

3. 数え上げの実装

余事象を考えると扱いやすくなる、と先ほど述べたポイントです。
X := \{ X_1,...X_L \} | (X_i =0,1)を、i桁目がX_iの2進数だと考えます。

例えばX = \{ 0,1,0,1 \} は1010です。
このケースだと、「2桁目で"0だけが出現している"or"1だけが出現している"ような選び方」 かつ「4桁目で"0だけが出現している"or"1だけが出現している"ような選び方」 の数がf(X)になります。

具体例を見てみましょう。
f(X)は下記の場合分けの合計になります。
・(2桁目,4桁目) = (0,0)
0000,0001,0100,0101 などからいくつか選ぶ方法
→　A_i \; and \; 1010 = 0000となるような数A_iからいくつか選ぶ

・(2桁目,4桁目) = (0,1)
1000,1001,1100,1101 などからいくつか選ぶ方法
→　A_i\; and \; 1010 = 0010となるような数A_iからいくつか選ぶ

・(2桁目,4桁目) = (1,0)
0010,0011,0110,0111 などからいくつか選ぶ方法
→　A_i\; and \; 1010 = 1000となるような数A_iからいくつか選ぶ

・(2桁目,4桁目) = (1,1)
1010,1011,1110,1111 などからいくつか選ぶ方法
→　A_i \; and \; 1010 = 1010となるような数A_iからいくつか選ぶ

このようにみると、Xをビットマスクと捉えて、
A_iをマスク処理した後に同値になるグループのサイズがわかれば、f(X)が計算できることがわかります。
（いくつか選ぶ＝1~K個選ぶ方法の数は二項係数の和で計算できます。）

よって、ビットマスクの値を一つ固定し、N個のA_iにそれぞれマスク処理をして、同値になるグループごとに選び方を計算する、という処理をすることでこの問題を解くことができます。

ビットマスクの数が2^ L、
マスク処理の回数がN
1回のマスク処理でL桁処理する（※）ので、
これらはO(2^ LLN) \approx O(2.4*10^ 8)で計算できます。
※　この処理はbit演算で速い、かつ桁数も少ないので致命的にはならなそう
※※　二項係数の和はO(N^ 2)で事前計算しておきます。（TLE1回）

n,k,s,t = map(int,input().split())
a = list(map(int,input().split()))

# 二項係数の事前計算
fac = [1] * (n+1)
for i in range(1,n+1):
    fac[i] = fac[i-1] * i
com = [[0] * (n+1) for _ in range(n+1)]
for i in range(n+1):
    for j in range(i+1):
        com[i][j] = fac[i] // (fac[j] * fac[i-j])

com_n = [0] * (n+1)
for i in range(n+1):
    for j in range(1,k+1):
        com_n[i] += com[i][j]

# 要らない数、桁の圧縮
use = [-1] * 18
for i in range(18):
    si = (s >> i) & 1
    ti = (t >> i) & 1
    if si == 0 and ti == 0:
        use[i] = 0
    elif si == 0 and ti == 1:
        use[i] = 2
    elif si == 1 and ti == 0:
        print(0)
        exit()
    else:
        use[i] = 1

a2 = []
for ai in a:
    new = 0
    for j in range(17,-1,-1):
        x = (ai >> j) & 1
        if use[j] == 0 and x == 1:
            break
        if use[j] == 1 and x == 0:
            break
        if use[j] == 2:
            new = new*2 + x
    else:
        a2.append(new)

# 包除原理で数え上げ
l = use.count(2) # 圧縮後の桁数
ans = 0
for mask in range(2**l):
    group_cnt = dict()
    for ai in a2:
        masked = ai & mask
        if masked in group_cnt:
            group_cnt[masked] += 1
        else:
            group_cnt[masked] = 1

    pc = bin(mask).count('1')    
    for j in group_cnt.values():
        ans += (-1) ** pc * com_n[j]

print(ans)

感想

余事象に言い換えるところが全く発想になく、3^ Lの世界で悩み続けていました。
もし言い換えができたとしても、包除原理できることに気づく、ビットマスクを使うとf(X)が数え上げられる、なども発想の飛躍があり厳しいです。

ちなみに、公式解説だと計算量O(2^ LN)となっており、Lを外せるらしいのですが、そこの工夫はよくわかっていません。
余事象に言い換えるところはmaspyさんのブログを参考にしました。