結果

問題 No.206 数の積集合を求めるクエリ
コンテスト
ユーザー かりあげクン
提出日時 2020-08-25 13:33:56
言語 Nim
(2.2.0)
結果
AC  
実行時間 400 ms / 7,000 ms
コード長 3,128 bytes
コンパイル時間 4,031 ms
コンパイル使用メモリ 66,328 KB
実行使用メモリ 73,480 KB
最終ジャッジ日時 2024-11-06 11:17:03
合計ジャッジ時間 15,136 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 28
権限があれば一括ダウンロードができます

ソースコード

diff # 

import strutils
import sequtils
import math
import lenientops

type Complex* = tuple[re: float, im: float]
proc complex*(x: float, y: float = 0.0): Complex {.inline.} =
  (x, y)
proc `+`*(a: Complex, b: Complex): Complex {.inline.} =
  (a.re + b.re, a.im + b.im)
proc `-`*(a: Complex, b: Complex): Complex {.inline.} =
  (a.re - b.re, a.im - b.im)
proc `*`*(a: Complex, b: Complex): Complex {.inline.} =
  (a.re * b.re - a.im * b.im, a.re * b.im + a.im * b.re)
proc `/`*(a: Complex, b: Complex): Complex {.inline.} =
  ( (a.re * b.re + a.im * b.im) / (b.re * b.re + b.im * b.im),
    (a.im * b.re - a.re * b.im) / (b.re * b.re + b.im * b.im) )
proc `*`*[T](a: Complex, k: T): Complex {.inline.} =
  (a.re * k, a.im * k)
proc `/`*[T](a: Complex, k: T): Complex {.inline.} =
  (a.re / k, a.im / k)
proc inv*(a: Complex): Complex {.inline.} =
  (a.re / a.re * a.re + a.im * a.im,
  -a.im / a.re * a.re + a.im * a.im)
proc `+=`*(a: var Complex, b: Complex) {.inline.} =
  a = a + b
proc `-=`*(a: var Complex, b: Complex) {.inline.} =
  a = a - b
proc `*=`*(a: var Complex, b: Complex) {.inline.} =
  a = a * b
proc `/=`*(a: var Complex, b: Complex) {.inline.} =
  a = a / b
proc `*=`*[T](a: var Complex, k: T) {.inline.} =
  a = a * k
proc `/=`*[T](a: var Complex, k: T) {.inline.} =
  a = a / k

proc dft*(F: var seq[Complex]): seq[Complex] =
  let N = F.len
  let mask = N - 1
  var tmp = newSeq[Complex](N)
  var i = N
  while i > 1:
    i = i shr 1
    let theta = 2 * PI * i / N
    let zeta = complex(cos(theta), sin(theta))
    var powZeta = complex(1.0)
    var j = 0
    while j < N:
      for k in 0 .. i - 1:
        tmp[j + k] = F[((j shl 1) and mask) + k] + powZeta * F[(((j shl 1) + i) and mask) + k]
      powZeta *= zeta
      j += i
    swap(F, tmp)
  return F

proc idft*(F: var seq[Complex]): seq[Complex] =
  let N = F.len
  let mask = N - 1
  var tmp = newSeq[Complex](N)
  var i = N
  while i > 1:
    i = i shr 1
    let theta = -1 * 2 * PI * i / N
    let zeta = complex(cos(theta), sin(theta))
    var powZeta = complex(1.0)
    var j = 0
    while j < N:
      for k in 0 .. i - 1:
        tmp[j + k] = F[((j shl 1) and mask) + k] + powZeta * F[(((j shl 1) + i) and mask) + k]
      powZeta *= zeta
      j += i
    swap(F, tmp)
  F.applyIt(it / N.float)
  return F

proc multiply*(A: seq[Complex], B: seq[Complex]): seq[int] =
  let N = nextPowerOfTwo(A.high + B.high + 1)
  var invA = A
  invA.setLen(N)
  invA = dft(invA)
  var invB = B
  invB.setLen(N)
  invB = dft(invB)
  var invF = newSeq[Complex](N)
  for i in 0 .. N - 1:
    invF[i] = invA[i] * invB[i]
  let F = idft(invF).mapIt(int(round(it.re)))
  return F

const FFT_max : int = 1 shl 18
var
  lmn     = stdin.readLine.split.map(parseInt)
  A       = stdin.readLine.split.map(parseInt)
  B       = stdin.readLine.split.map(parseInt)
  Q       = stdin.readLine.parseInt
  (l,m,n) = (lmn[0], lmn[1], lmn[2])
  A_poly  = newSeq[Complex](FFT_max)
  B_poly  = newSeq[Complex](FFT_max)
for i in 0..<l:
  A_poly[A[i] - 1] = (1.0, 0.0)
for i in 0..<m:
  B_poly[n - B[i]] = (1.0, 0.0)

var C = multiply(A_poly, B_poly)

for i in 0..<Q:
  echo C[n-1+i]
0