結果

問題 No.206 数の積集合を求めるクエリ
ユーザー qwewe
提出日時 2025-05-14 13:27:10
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 73 ms / 7,000 ms
コード長 3,805 bytes
コンパイル時間 938 ms
コンパイル使用メモリ 101,808 KB
実行使用メモリ 16,520 KB
最終ジャッジ日時 2025-05-14 13:27:40
合計ジャッジ時間 3,311 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 28
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <vector>
#include <complex>
#include <cmath>
#include <algorithm> // For std::swap

using Complex = std::complex<double>;
const double PI = acos(-1.0); // M_PI is not standard C++, acos(-1.0) is a common way

// Iterative FFT function
void fft(std::vector<Complex>& a, bool invert) {
    int n = a.size();
    if (n == 1) return;

    // Bit-reversal permutation
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1)
            j ^= bit;
        j ^= bit;
        if (i < j)
            std::swap(a[i], a[j]);
    }

    // Iterative Cooley-Tukey
    for (int len = 2; len <= n; len <<= 1) {
        double ang = 2 * PI / len * (invert ? -1 : 1);
        Complex wlen(cos(ang), sin(ang));
        for (int i = 0; i < n; i += len) {
            Complex w(1);
            for (int j = 0; j < len / 2; j++) {
                Complex u = a[i + j];
                Complex v = a[i + j + len / 2] * w;
                a[i + j] = u + v;
                a[i + j + len / 2] = u - v;
                w *= wlen;
            }
        }
    }

    if (invert) {
        for (int i = 0; i < n; i++)
            a[i] /= n; // Scale by 1/n for inverse FFT
    }
}


int main() {
    std::ios_base::sync_with_stdio(false); // Faster I/O
    std::cin.tie(NULL); // Untie cin from cout

    int L, M, N_val;
    std::cin >> L >> M >> N_val;

    std::vector<int> A_list(L);
    for (int i = 0; i < L; ++i) {
        std::cin >> A_list[i];
    }

    std::vector<int> B_list(M);
    for (int i = 0; i < M; ++i) {
        std::cin >> B_list[i];
    }

    int Q_val;
    std::cin >> Q_val;

    int s_fft = 1;
    // Smallest N_val is 1 due to constraints 1 <= N.
    // So 2*N_val is at least 2.
    // s_fft must be a power of 2 and s_fft >= 2*N_val.
    while (s_fft < 2 * N_val) { 
        s_fft <<= 1;
    }
    // If N_val = 1, 2*N_val = 2. s_fft starts at 1.
    // Loop: 1 < 2 is true, s_fft becomes 2. 2 < 2 is false. s_fft = 2. Correct.
    // If 2*N_val is 0 (e.g. N_val=0, not allowed by constraints), loop condition false, s_fft=1.

    std::vector<Complex> pA(s_fft, {0.0, 0.0});
    for (int val_a : A_list) {
        // A[i] are values from 1 to N_val. These become indices.
        pA[val_a] = {1.0, 0.0};
    }

    std::vector<Complex> pB_rev(s_fft, {0.0, 0.0});
    for (int val_b : B_list) {
        // B[i] are values from 1 to N_val.
        // Index N_val - val_b (ranges from 0 to N_val-1).
        int idx_b_rev = N_val - val_b;
        pB_rev[idx_b_rev] = {1.0, 0.0};
    }
    
    fft(pA, false); 
    fft(pB_rev, false);

    std::vector<Complex> pC_fft_coeffs(s_fft);
    for (int i = 0; i < s_fft; ++i) {
        pC_fft_coeffs[i] = pA[i] * pB_rev[i];
    }

    fft(pC_fft_coeffs, true); // Inverse FFT

    for (int v_query = 0; v_query < Q_val; ++v_query) {
        // We need coefficient for a-b = v_query.
        // This is poly_C[v_query + N_val].
        int k_index = v_query + N_val;
        
        long long ans = 0;
        // k_index is always >= N_val.
        // Max k_index is (Q-1)+N_val. Since Q <= N_val, max k_index <= (N_val-1)+N_val = 2*N_val-1.
        // This is within s_fft bounds because s_fft >= 2*N_val.
        // So, k_index < s_fft is guaranteed.
        // A check `if (k_index >=0 && k_index < s_fft)` is technically redundant for k_index < s_fft
        // if N_val > 0. For k_index >= 0, since N_val >= 1 and v_query >= 0, k_index >= 1.
        // So the check can be simplified, or kept for safety.
        
        ans = static_cast<long long>(pC_fft_coeffs[k_index].real() + 0.5);
        
        // Counts must be non-negative. Floating point artifacts might make it slightly negative.
        if (ans < 0) ans = 0; 
        std::cout << ans << "\n";
    }

    return 0;
}
0