高速フーリエ変換の実装

require 'complex'

# 原始N乗根
def primitiveNthRoot(n, k=1)
  t = 2 * Math::PI * k / n
  Complex(Math.cos(t), Math.sin(t))
end

# 奇数？
def odd(x)
  (x & 1) != 0
end

# バタフライ
def f(a)
  n = a.length
  if n == 1
    a
  elsif odd(n)
    raise "not even"
  else
    n2 = n / 2
    ae = (0...n2).map {|l| a[l*2]}  # 偶数成分
    ao = (0...n2).map {|l| a[l*2 + 1]}  # 奇数成分

    p = f(ae)
    q = f(ao)

    b = Array.new(n)
    (0...n2).each do |k|
      wkN = primitiveNthRoot(n, k)
      b[k]      = p[k] + wkN * q[k]
      b[k + n2] = p[k] - wkN * q[k]
    end
    b
  end
end

# 高速フーリエ変換
def fft(fs)
  n = fs.length
  b = f(fs)
  b.map {|be| be.conjugate / n}
end

if $0 == __FILE__
  v = nil
  if ARGV.empty?
    # 固定の信号
    v = [1,7,3,2,0,5,0,8]
  else
    # 指定の長さの信号をランダムで生成
    n = ARGV.shift.to_i
    v = (0...n).map {|i| rand * 2 - 1}
  end

  # 高速フーリエ変換
  r = fft(v)
  p r

  #=== 検算

  # 逆フーリエ変換
  def ifft(cs)
    b = f(cs)
    b.map {|c| c.real}
  end
  s = ifft(r)

  # 元の値との誤差
  p v.vsub(s).vlength
end