東京海上日動プログラミングコンテスト2020のE問題O(rand)の包除原理を用いた解法のWAの原因を教えていただきたいです。

Question

実現したいこと

東京海上日動プログラミングコンテスト2020のE問題O(rand)で包除原理を用いて解いたのですが、サンプルケースは全てACになるもののテストケースのほとんどでWAになってしまいました。
解説や他の方の解答などを読んでも方針は間違っていないと思いオーバフローなどもないように思われるので、原因の究明をお願いしたいです。
また、デバッグしていただきやすいようにコメントアウトでやりたいことを書いてあります。よろしくお願いします。

ソースコード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

//マクロ
//forループ関係
//引数は、(ループ内変数,動く範囲)か(ループ内変数,始めの数,終わりの数)、のどちらか
#define REP(i,n) for(ll i=0;i<(ll)(n);i++)
#define FOR(i,a,b) for(ll i=a;i<=(ll)(b);i++)
//xにはvectorなどのコンテナ
#define ALL(x) (x).begin(),(x).end() //sortなどの引数を省略したい
#define SIZE(x) ((ll)(x).size()) //sizeをsize_tからllに直しておく
//定数
#define INF 1000000000000 //10^12:極めて大きい値,∞
#define MOD 1000000007 //10^9+7:合同式の法
#define MAXR 100 //10^5:配列の最大のrange(素数列挙などで使用)
//略記
#define PB push_back //vectorヘの挿入
#define MP make_pair //pairのコンストラクタ
#define F first //pairの一つ目の要素
#define S second //pairの二つ目の要素


signed main(){
    //入力の高速化用のコード
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    ll n,k,s,t;cin >> n >> k >> s >> t;
    vector<ll> b(n);REP(i,n)cin >> b[i];
    vector<ll> a;
    //i bit目を見て、s_i=1.t_i=1の時は全ての数のi bit目が1→当てはまらないのは除く
    //i bit目を見て、s_i=0.t_i=0の時は全ての数のi bit目が0→当てはまらないのは除く
    //i bit目を見て、s_i=1,t_i=0の時は0を出力
    //i bit目を見て、s_i=0,t_i=1の時は候補に入れる
    REP(i,n){
        bool f=true;
        REP(j,18){
            bool s_,t_,b_;s_=(s>>j)&1;t_=(t>>j)&1;b_=(b[i]>>j)&1;
            if(s_ and t_){
                if(!(b_)){
                    f=false;break;
                }
            }else if(!(s_) and !(t_)){
                if(b_){
                    f=false;break;
                }
            }else if(s_ and !(t_)){
                cout<<0<<endl;return 0;
            }
        }
        if(f)a.PB(b[i]);
    }
    n=SIZE(a);
    k=min(n,k);
    //以下ではそれぞれのbitは0と1が少なくとも一つなければならない
    //余事象は、全ての数で1だけまたは0だけのbitが少なくとも一つある場合
    //少なくとも一つ→包除原理
    //それぞれのbitで"0だけになる集合"と"1だけになる集合"を用意すると、3^L(L=18)になる→O(N×3^L)
    //0だけor1だけなので分割してく
    //チェックするのは結局N個くらいで済む→O(N×2^L×L)
    
    //以下、集合はbitで表現

    //まずは、それぞれのbitで集合を分割(そのbitが1か0かで二つに分割→pairで保存)
    vector<pair<ll,ll>> bits;
    REP(i,18){
        ll f1=0;ll f2=0;
        REP(j,n){
            if((a[j]>>i)&1){
                f1+=(1<<j);
            }else{
                f2+=(1<<j);
            }
        }
        if(f1 and f2)bits.PB(MP(f1,f2));
    }
    ll bits_len=SIZE(bits);

    //集合分割したのでさらに分割を行う(部分集合全列挙したいのでbit全探索)
    //余事象が何通りあるかをans_subに格納
    ll ans_sub=0;

    //ここで答えでの組み合わせの前計算しておく
    //prior_calc[i][j]：iC1+iC2+…+iCjの答え(i:1~n,j:1~k)
    //0-indexed
    vector<vector<ll>> pcalc(n+1,vector<ll>(k+1,0));
    pcalc[0][0]=1;
    FOR(i,1,n){
        pcalc[i][0]=1;
        FOR(j,1,k){
            pcalc[i][j]=pcalc[i-1][j-1]+pcalc[i-1][j];
        }
    }
    REP(i,n+1){
        pcalc[i][0]=0;
        REP(j,k){
            pcalc[i][j+1]+=pcalc[i][j];
        }
    }

    //bit全探索でありうる積集合のパターンを全て考える
    //2^L
    //all_setsにはそこまでで分けられた集合を集合ごとに保存(bitでどの要素が含まれるか表す)
    deque<ll> all_sets;
    FOR(i,1,(1<<bits_len)-1){
        //まず全ての要素が同じ集合に含まれているものとする
        all_sets.PB((1<<n)-1);
        //L
        //包除原理において、積集合で考える集合の個数の偶奇が知りたいので
        ll odd_even=0;
        REP(j,bits_len){
            //そのbitを積集合を考えるbitとして考慮に入れる時
            if((i>>j)&1){
                odd_even+=1;
                ll s=SIZE(all_sets);
                //n
                //今分けられている集合それぞれをさらに(そのbitが0か1かで)分割する
                REP(_,s){
                    //積集合をとりたい集合を書き出す
                    ll ins=*(all_sets.begin());
                    all_sets.pop_front();
                    //insはf1とf2へと分割される
                    ll f1=ins&(bits[j].F);ll f2=ins&(bits[j].S);
                    //f1,f2が0の時は空集合なので0通りより無視
                    if(f1){all_sets.PB(f1);}
                    if(f2){all_sets.PB(f2);}
                }
            }
        }
        //積集合のパターンを決めた下で完全に集合を分割できた
        //それぞれの集合に対して1以上K以下の数を選ぶ(pcalcを使う)
        ll s=SIZE(all_sets);
        //n
        REP(_,s){
            ll now=*(all_sets.begin());
            all_sets.pop_front();
            //立ってるbit数を数える(注目する集合の要素数を数える)
            ll s_sub=__builtin_popcount(now);
            //前計算で求めたpcalcから求める
            //ここで、包除原理より+-が決まる
            if(odd_even%2){
                ans_sub+=pcalc[s_sub][min(k,s_sub)];
            }else{
                ans_sub-=pcalc[s_sub][min(k,s_sub)];
            }
        }
    }
    //ans_subは余事象のパターン数で、全体のパターン数からその余事象のパターン数を引く
    cout<< pcalc[n][k]-ans_sub <<endl;
}

Answer 1

        if(f1 and f2)bits.PB(MP(f1,f2));

どちらかが0 かつSのビットが0でTのビットが1 の場合、条件を満たす組み合わせがないということなので、push_backしないだけじゃなくて0を出力して終了するような分岐にすべきです。

あとは型の問題で
シフト演算は 1LL<<i というように左辺をlong long 型にしないとlong long 型にはなりません(オーバーフローします)。
__builtin_popcountも64ビット版は　__builtin_popcountllという名前のようです。

これでACできるかどうかはわかりませんが......

---

もうAC取れたということで蛇足かもしれませんが

どちらかが0 かつどちらかが0 かつSのビットが0でTのビットが1 の場合

今のコードで短絡評価してるのはSのビットが1でTのビットが0の場合なので、ここは評価されていません。
これが問題になるのは下のようなケースです

2 1 0 3
2 4

明らかにTが3になる組み合わせは存在しませんが、今のコードだと1が出力されます。

Comments

[退会済みユーザー]

> どちらかが0かつSのビットが0でTのビットが1の場合、条件を満たす組み合わせがないということなので、push_backしないだけじゃなくて0を出力して終了するような分岐にすべきです。

この記述についてなのですが、Sのビットが0でTのビットが1の場合は数行前で0を出力して終了しているので、ここでは考慮していないです。自分の書き方がわかりにくかったと思います。ご指摘ありがとうございます。

`1LL<<i`と`__builtin_popcountll`の二つのご指摘の点を直したところ無事ACできました。テストケースでSとTが大きい時にWAだったのでオーバフローだと思ったのですが、勉強不足で上記二つの点にたどり着くことができませんでした。

ご丁寧に回答していただきACできたので、ベストアンサーとさせていただきます。非常に助かりました。ありがとうございました。

[退会済みユーザー]

蛇足についてですが、"どちらかが0かつSのビットが0でTのビットが1の場合"を"Sのビットが1でTのビットが0の場合"と読み間違っていました。
ご指摘の通り上記の場合では自分のコードでは間違った挙動を示すので書き直しました。ありがとうございました。