Skip to content

Instantly share code, notes, and snippets.

@yuppe19

yuppe19/dice.py

Last active September 10, 2015 01:04
Show Gist options
  • Save yuppe19/9fd4736d99ede5b663f4 to your computer and use it in GitHub Desktop.
Save yuppe19/9fd4736d99ede5b663f4 to your computer and use it in GitHub Desktop.
Download ZIP
(ΦωΦ)＜サイコロをN個振ったときの和 (高速フーリエ変換バージョン)
Raw
dice.py
#!/usr/bin/python
# -*- coding: utf-8 -*-
#http://www.slideshare.net/chokudai/fft-49066791
# (ΦωΦ)＜
# まだ精度が足りないのか、
# [114, 586]の範囲(要するにほとんど)で答えが違います。
from mpmath import mp, mpc, sin, cos, pi, power
mp.prec = 7777
def pow_2_at_least(x):
n = 1
while n < x:
n *= 2
return n
def convolution(g, nn):
Lg = len(g)
n = pow_2_at_least(Lg * nn)
g = g + [mpc(0, 0)] * (n - Lg)
gf = fft(g, n)
ff = [power(gf[i], nn) for i in xrange(n)]
return ifft(ff, n)
def _fft(v, n, theta):
if n == 1:
return v
m = n / 2
v0 = [v[i] for i in xrange(0, n, 2)]
v1 = [v[i] for i in xrange(1, n, 2)]
v0 = _fft(v0, m, theta)
v1 = _fft(v1, m, theta)
zeta = mpc(cos(theta/n), sin(theta/n))
pow_zeta = mpc(1, 0)
for i in xrange(n):
v[i] = v0[i % m] + pow_zeta * v1[i % m]
pow_zeta *= zeta
return v
fft = lambda v, n: _fft(v, n, 2*pi)
ifft = lambda v, n: map(lambda x: x/n, _fft(v, n, -2*pi))
K = 6
N = 100
for idx, e in enumerate(convolution([mpc(0, 0)] + [mpc(1, 0)] * K, N)[N: N*K+1]):
print idx+N, int(round(e.real))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment