#27312 (C++11) No.206 数の積集合を求めるクエリ

提出ソース

結果

問題	No.206 数の積集合を求めるクエリ
ユーザー	Min_25
提出日時	2015-05-16 00:59:06
言語	C++11 (gcc 11.4.0)
結果	AC
実行時間	38 ms / 7,000 ms
コード長	8,423 bytes
コンパイル時間	632 ms
コンパイル使用メモリ	70,056 KB
実行使用メモリ	7,808 KB
最終ジャッジ日時	2024-07-06 04:37:57
合計ジャッジ時間	2,528 ms
ジャッジサーバーID （参考情報）	judge2 / judge1

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,812 KB
testcase_01	AC	1 ms 6,944 KB
testcase_02	AC	2 ms 6,944 KB
testcase_03	AC	2 ms 6,940 KB
testcase_04	AC	1 ms 6,944 KB
testcase_05	AC	2 ms 6,940 KB
testcase_06	AC	2 ms 6,940 KB
testcase_07	AC	2 ms 6,944 KB
testcase_08	AC	2 ms 6,940 KB
testcase_09	AC	2 ms 6,940 KB
testcase_10	AC	2 ms 6,940 KB
testcase_11	AC	1 ms 6,940 KB
testcase_12	AC	3 ms 6,940 KB
testcase_13	AC	2 ms 6,940 KB
testcase_14	AC	2 ms 6,944 KB
testcase_15	AC	3 ms 6,940 KB
testcase_16	AC	3 ms 6,940 KB
testcase_17	AC	36 ms 7,680 KB
testcase_18	AC	34 ms 7,808 KB
testcase_19	AC	34 ms 7,680 KB
testcase_20	AC	34 ms 7,680 KB
testcase_21	AC	34 ms 7,680 KB
testcase_22	AC	35 ms 7,680 KB
testcase_23	AC	35 ms 7,680 KB
testcase_24	AC	38 ms 7,680 KB
testcase_25	AC	37 ms 7,680 KB
testcase_26	AC	36 ms 7,680 KB
testcase_27	AC	35 ms 7,680 KB
testcase_28	AC	37 ms 7,680 KB
testcase_29	AC	37 ms 7,680 KB
testcase_30	AC	37 ms 7,680 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include <cstdio>
#include <cmath>
#include <iostream>
#include <complex>

#define getchar getchar_unlocked
#define putchar putchar_unlocked

using namespace std;

typedef long long int64;
typedef long long unsigned uint64;
typedef long double float80;
typedef unsigned short uint16;
typedef unsigned uint;
typedef unsigned char uint8;

uint get_uint() {
  uint n;
  int c;
  while( (c = getchar()) < '0') {
    ;
  }
  n = c - '0';
  while( (c = getchar()) >= '0') {
    n = n * 10 + (c - '0');
  }
  return n;
}

void put_uint(uint n) {
  uint8 stack[30];
  int top = 0;
  do {
    stack[top++] = n % 10 + '0';
    n /= 10;
  } while(n != 0);
  while(top > 0) {
    putchar(stack[--top]);
  }
  putchar('\n');
}

typedef double real_t;

template <typename real_t>
class Complex {
public:
  Complex() {}
  Complex(const real_t a) : real_(a), imag_(0) {}
  Complex(const real_t a, const real_t b) : real_(a), imag_(b) {}
  Complex(const Complex& t) : real_(t.real_), imag_(t.imag_) {}
  Complex operator+(Complex rhs) const {
    Complex ret = Complex(*this);
    ret.real_ += rhs.real_;
    ret.imag_ += rhs.imag_;
    return ret;
  }
  Complex operator-(Complex rhs) const {
    Complex ret = Complex(*this);
    ret.real_ -= rhs.real_;
    ret.imag_ -= rhs.imag_;
    return ret;
  }
  Complex operator*(Complex rhs) const {
    Complex ret = Complex(
      real_ * rhs.real_ - imag_ * rhs.imag_,
      real_ * rhs.imag_ + imag_ * rhs.real_
    );
    return ret;
  }
  Complex operator*(const real_t rhs) const {
    Complex ret = Complex(*this);
    ret.real_ *= rhs;
    ret.imag_ *= rhs;
    return ret;
  }
  Complex operator+=(const Complex rhs) {
    return *this = *this + rhs;
  }
  Complex operator-=(const Complex rhs) {
    return *this = *this - rhs;
  }
  Complex operator*=(const Complex rhs) {
    return *this = *this * rhs;
  }
  Complex operator*=(const real_t rhs) {
    return *this = *this * rhs;
  }
  real_t real() const {
    return real_;
  }
  real_t imag() const {
    return imag_;
  }
private:
  real_t real_, imag_;
};

typedef Complex<real_t> complex_t;

class FFT {
private:
  static constexpr real_t PI = acosl(-1);
  FFT() {}

public:
  static void convolute_reals(real_t* A, const uint size_a, real_t* B, const uint size_b) {
    uint size = std::max(size_a, size_b);
    uint ldn = 0;
    uint fft_size = 1;
    while (fft_size < size) {
      ++ldn;
      fft_size <<= 1;
    }
    ++ldn;
    fft_size <<= 1;

    std::fill(A + size_a, A + fft_size, real_t(0));
    std::fill(B + size_b, B + fft_size, real_t(0));

    real_to_complex_fft(A, ldn);
    real_to_complex_fft(B, ldn);

    for (uint i = 1; i < fft_size / 2; ++i) {
      complex_t a = complex_t(A[2 * i], A[2 * i + 1]);
      complex_t b = complex_t(B[2 * i], B[2 * i + 1]);
      a *= b;
      A[2 * i + 0] = a.real();
      A[2 * i + 1] = a.imag();
    }
    A[0] *= B[0];
    if (fft_size > 1) {
      A[1] *= B[1];
    }

    complex_to_real_fft(A, ldn);
    real_t inv = 1. / fft_size;
    for (uint i = 0; i < fft_size; ++i) {
      A[i] *= inv;
    }
  }

  static void convolute(complex_t* A, const uint size_a, complex_t* B, const uint size_b) {
    uint size = std::max(size_a, size_b);
    uint ldn = 0;
    uint fft_size = 1;
    while (fft_size < size) {
      ++ldn;
      fft_size <<= 1;
    }
    ++ldn;
    fft_size <<= 1;

    std::fill(A + size_a, A + fft_size, complex_t(0));
    std::fill(B + size_b, B + fft_size, complex_t(0));

    fft_dit4(A, ldn, 1);
    fft_dit4(B, ldn, 1);
    for (uint i = 0; i < fft_size; ++i) {
      A[i] *= B[i];
    }
    fft_dit4(A, ldn, -1);

    real_t inv = real_t(1) / fft_size;
    for (uint i = 0; i < fft_size; ++i) {
      A[i] *= inv;
    }
  }

  static void fft_dit4(complex_t *f, uint ldn, int sign) {
    const uint LX = 2;
    const uint n = 1u << ldn;

    revbin_permute(f, n);

    if (ldn & 1) {
      for (uint r = 0; r < n; r += 2) {
        sumdiff(f[r], f[r + 1]);
      }
    }
    uint ldm = LX + (ldn & 1);

    for (; ldm <= ldn; ldm += LX) {
      uint m = 1u << ldm;
      uint m4 = m >> LX;
      real_t ph0 = (2.0 * PI * sign) / m;

      for (uint j = 0; j < m4; ++j) {
        real_t phi = j * ph0;
        real_t s, c;
        sin_cos(phi, s, c);
        complex_t e  = complex_t(c, s);
        complex_t e2 = e * e;
        complex_t e3 = e2 * e;

        for (uint r = 0; r < n; r += m) {
          const uint i = j + r;
          complex_t a0 = f[i + m4 * 0];
          complex_t a2 = f[i + m4 * 1] * e2;
          complex_t a1 = f[i + m4 * 2] * e;
          complex_t a3 = f[i + m4 * 3] * e3;

          complex_t t02 = a0 + a2;
          complex_t t13 = a1 + a3;

          f[i + m4 * 0] = t02 + t13;
          f[i + m4 * 2] = t02 - t13;

          t02 = a0 - a2;
          t13 = (a1 - a3) * complex_t(0, sign);

          f[i + m4 * 1] = t02 + t13;
          f[i + m4 * 3] = t02 - t13;
        }
      }
    }
  }

private:
  public: static void sin_cos(const real_t rad, real_t& s, real_t& c) {
    s = sin(rad);
    c = sqrt(1 - s * s);
  }

  static void real_to_complex_fft(real_t* f, uint ldn) {
    if (ldn == 0) {
      return;
    }

    fft_dit4((complex_t*)f, ldn - 1, 1);

    const uint n = 1u << ldn;
    const uint mh = n / 2, n4 = n / 4;
    const real_t phi0 = PI / mh;

    // before:
    // - F(a + bi) = c_s - d_a + i * (c_a + d_s)
    //
    // after:
    // - F(a + bz) =      c_s + d_s * cos - d_a * sin
    //             + i * (c_a + d_s * sin + d_a * cos)
    for (uint i = 1; i < n4; ++i) {
      uint i1 = 2 * i;
      uint i3 = n - i1;

      real_t r_k = f[i1 + 0], r_nk = f[i3 + 0];
      sumdiff05(r_nk, r_k); // r_nk = c_s, r_k = d_a

      real_t i_k = f[i1 + 1], i_nk = f[i3 + 1];
      sumdiff05(i_k, i_nk); // i_k = d_s, i_nk = c_a

      real_t phi = i * phi0;
      real_t s, c;
      sin_cos(phi, s, c);

      // i_k = d_s * cos - d_a * sin
      // r_k = d_s * sin + d_a * cos
      cmult(c, s, i_k, r_k);
      f[i1 + 0] = r_nk + i_k;
      f[i3 + 0] = r_nk - i_k;

      f[i1 + 1] =  i_nk + r_k;
      f[i3 + 1] = -i_nk + r_k;
    }
    sumdiff(f[0], f[1]);
  }

  static void complex_to_real_fft(real_t* f, uint ldn) {
    if (ldn == 0) {
      return;
    }

    const uint n = 1u << ldn;
    const uint nh = n / 2, n4 = n / 4;
    const real_t phi0 = -PI / nh;

    // before:
    // - F(a - bi) = c_s + d_a + i * (c_a - d_s)
    //
    // after:
    // - F(a + bz^{-1}) =      c_s + d_s * cos + d_a * sin
    //                  + i * (c_a - d_s * sin + d_a * cos)
    //
    for (uint i = 1; i < n4; ++i) {
      uint i1 = 2 * i;
      uint i3 = n - i1;

      real_t r_k = f[i1 + 0], r_nk = f[i3 + 0];
      sumdiff(r_k, r_nk); // r_k = c_s, r_nk = d_a

      real_t i_k = -f[i1 + 1], i_nk = f[i3 + 1];
      sumdiff(i_k, i_nk); // i_k = -c_a, i_nk = d_s

      real_t phi = i * phi0;
      real_t s, c;
      sin_cos(phi, s, c); // cos, -sin

      // i_nk =  d_s * cos + d_a * sin
      // r_nk = -d_s * sin + d_a * cos
      cmult(c, s, i_nk, r_nk);

      f[i1 + 0] = r_k + i_nk;
      f[i3 + 0] = r_k - i_nk;

      f[i1 + 1] = -i_k + r_nk;
      f[i3 + 1] =  i_k + r_nk;
    }
    sumdiff(f[0], f[1]);

    if (nh >= 2) {
      f[nh] *= 2.0;
      f[nh + 1] *= 2.0;
    }

    fft_dit4((complex_t*)f, ldn - 1, -1);
  }

  static inline void cmult(real_t c, real_t s, real_t& u, real_t& v) {
    real_t t = u * s + v * c;
    u *= c;
    u -= v * s;
    v = t;
  }

  template <typename T>
  static inline void sumdiff05(T& a, T& b) {
    T t = (a - b) * 0.5;
    a += b;
    a *= 0.5;
    b = t;
  }

  template <typename T>
  static inline void sumdiff(T& a, T& b) {
    T t = a - b;
    a += b;
    b = t;
  }

  template <typename T>
  static void revbin_permute(T* A, uint n) {
    if(n <= 2) {
      return;
    }
    uint r = 0;
    uint nh = n >> 1;
    for(uint x = 1; x < n; ++x) {
      uint h = nh;
      while(! ((r ^= h) & h)) {
        h >>= 1;
      }
      if(r > x) {
        std::swap(A[x], A[r]);
      }
    }
  }
};

real_t A[1 << 18];
real_t B[1 << 18];

int main() {
  uint L = get_uint();
  uint M = get_uint();
  uint N = get_uint();

  for (uint i = 0; i < L; ++i) {
    A[get_uint()] += 1;
  }
  for (uint i = 0; i < M; ++i) {
    B[N - get_uint()] += 1;
  }

  // 2n real => n complex FFT
  FFT::convolute_reals(A, N + 1, B, N + 1);

  uint Q = get_uint();

  for (uint i = N; i < N + Q; ++i) {
    put_uint(uint(A[i]+ 0.5));
  }
  return 0;
}

yukicoder

結果

テストケース

ソースコード