mayoko’s diary

プロコンとかいろいろ。

yukicoder No.206 数の積集合を求めるクエリ

なんか久しぶりに記事書いてる気がする…
この問題はあとでFFT使って解いてみたいです(まだFFTをよく理解していないのでできませんが…)

解法

基本は解説(http://yukicoder.me/problems/440/editorial)に書いてあるとおりで,問題は実装方法ですね。bitsetというものを初めて使いました。
例えばAの配列でどの数字が出てきているのかというのをbaというbitsetの配列に保持しておくとすると,

ba[0]に0〜63,ba[1]に64〜127,...というように64おきにどの数字が出てきているのか,ということを管理しています。bb(Bの配列においてどの数字が出てくるかを管理するbitsetの配列)についてもどの数字が出てきているのかを管理しておけば,ba[i]&bb[i]をやることで「どの数字が共通しているのか」がわかります。

また,B[]の値を1ずつずらしていくのはシフト演算で実行することができます。

以下ソースコード

const int MAXN = 1e5+5;
int A[MAXN];
int B[MAXN];
bitset<64> ba[2000], bb[2000];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int L, M, N;
    cin >> L >> M >> N;
    for (int i = 0; i < L; i++) {
        cin >> A[i];
        ba[A[i]/64].set(A[i]%64);
    }
    for (int i = 0; i < M; i++) {
        cin >> B[i];
        bb[B[i]/64].set(B[i]%64);
    }
    int Q;
    cin >> Q;
    while (Q--) {
        int ans = 0;
        for (int i = 0; i < 2000; i++) {
            ans += (ba[i]&bb[i]).count();
        }
        cout << ans << endl;
        for (int i = 1999; i > 0; i--) {
            bb[i] <<= 1;
            if (bb[i-1].test(63)) bb[i].set(0);
        }
        bb[0] <<= 1;
    }
    return 0;
}