numpy.poly1dと母関数とサイコロ

　「部屋とYシャツと私」みたいなタイトルになってしまいました。

　サイコロを振ったときの確率・統計量を母関数を使って計算するテクニックを知ったのでその紹介をします。
　母関数を扱う際に多項式の計算ライブラリがあると嬉しいです。今回はPythonのnumpy.poly1dクラスを使います。

numpy.poly1dクラスの使用例

多項式の加算、乗算、べき乗、根の計算、微分、特定の点での評価などが出来ます。

from numpy import poly1d as poly

f = poly([1,2,3])    # f(x) = x^2  + 2x + 3
g = poly([2, 1])     # g(x) = 2x + 1

# addition, multiplication and power
print f + g          
print f * g
print g ** 5

# roots of (x^2 - x - 1)
print poly([1, -1, -1]).r

# derivatives
print f.deriv()     # 1st derivative of f 
print f.deriv(2)    # 2nd derivative of f

# value at a value
print f(1)          # f(1) = 1^2 + 2 * 1 + 3

以下が実行結果です。指数の部分とそれ以外の部分を２行に分けて、いい感じに出力してくれます。

kenjih$ python sample.py 
   2
1 x + 4 x + 4
   3     2
2 x + 5 x + 8 x + 3
    5      4      3      2
32 x + 80 x + 80 x + 40 x + 10 x + 1
[ 1.61803399 -0.61803399]
 
2 x + 2
 
2
6

母関数とは？

応用に入る前に母関数について簡単に説明しておきます。

母関数とは、数列{a_n}に対して
 f(x) = ∑ a_n xⁿ
と定義される多項式のことです。

f(x)のxⁿの係数がa_nと等しいというのがポイントで、うまく利用すると「分割数の計算」や「漸化式の一般項の計算」が出来たりします。

サイコロの統計量と母関数

6つの面を持つサイコロを振ります。
各面が出る確率は等しいとします。

f(x) = (1/6) x + (1/6) x² + (1/6) x³ + (1/6) x⁴ + (1/6) x⁵ + (1/6) x⁶

という関数を定義します。xⁿの係数は「サイコロを振った時にnが出る確率」と同じです。

f(x)を1回微分してみます。

f '(x) = (1/6) + (1/6) 2x + (1/6) 3x² + (1/6) 4x³ + (1/6) 5x⁴ + (1/6) 6x⁵

1回微分することで、xⁿの指数部のnが係数のところに出てきます。これで何が嬉しいかというと、

f '(1) = (1/6) + (1/6) * 2 + (1/6) * 3 + (1/6) * 4 + (1/6) * 5 + (1/6) * 6

となり、x=1を代入すると平均値を計算することが出来ます。

次にf(x)を2回微分してみます。

f ''(x) = 0 + (1/6) * 2 + (1/6) 6x + (1/6) 12x² + (1/6) 20x³ + (1/6) 30x⁴

2回微分するxⁿの係数はn(n-1)になります。1回微分して出てきたnと足すと、n²が出てきます。
つまり、

f '(1) + f ''(1) = (1/6) + (1/6) * 2² + (1/6) * 3² + (1/6) * 4² + (1/6) * 5² + (1/6) * 6²

となり、二乗の平均が出せます。よって、

f '(1) + f ''(1) - (f '(1))²

で「二乗の平均 - 平均の二乗」、つまり、分散を計算することが出来ます。

ここまではサイコロを1回ふったときの話を考えました。
次はサイコロを複数回ふった場合、母関数がどうなるかを考えてみます。

f(x)の2乗

((1/6) x + (1/6) x² + (1/6) x³ + (1/6) x⁴ + (1/6) x⁵ + (1/6) x⁶) ²

を考えてみます。2乗を展開すると、x⁴ の係数は3/36になることが分かります。
これは、

(1/6) x¹ と (1/6) x³の積
(1/6) x² と (1/6) x²の積
(1/6) x³ と (1/6) x¹の積

の和です。

次のように言い換えることも出来ます。

1回目の試行で1が出て、2回目の試行で3が出た確率
1回目の試行で2が出て、2回目の試行で2が出た確率

1回目の試行で3が出て、2回目の試行で1が出た確率

の和。これは「サイコロを2回振った時の目の和が4になる確率」です。

よってf(x)の2乗を計算してxⁿの係数を見ることで、「サイコロを2回投げたときに目の和がnになる確率」を求めることが出来ます。

さらに一般化すると、f(x)のm乗を計算することで、「サイコロをm回投げた時の目の和の分布」を計算できるということになります。また、f(x)のm乗の微分を行うことで、平均や分散などを計算することができます。

例題

6面のサイコロがある。
サイコロには細工がされていて1,2,3,4,5,6が出る確率が1:2:3:4:5:6となっている。
サイコロを30回振った時に出た目の合計が120以上になる確率を求めよ。

まず、普通に動的計画法で解いてみます。

#include <iostream>

using namespace std;

const int N = 30;  // 30 iterations
const int K = 6;   // 6-sided dice

double dp[N+1][N*K+1];

int sumTo(int k) {
    return k * (k + 1) / 2;
}

int main(int argc, char **argv) {
    
    dp[0][0] = 1.0;
    
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N*K; j++) {
            if (dp[i][j] == 0)
                continue;
            for (int k = 1; k <= K; k++) {
                dp[i+1][j+k] += dp[i][j] * k / sumTo(K);
            }
        }
    }

    double ret = 0;
    for (int i = 120; i <= N*K; i++)
        ret += dp[N][i];

    cout << ret << endl;

    return 0;
}

答えは、0.898544となりました。

次に、母関数を用いた解法。

from numpy import poly1d as poly

dice = range(7)
total = sum(dice)
prob = [1. * x / total for x in dice]

f = poly(prob[::-1])
f = f ** 30

p = f.c[::-1]
print sum(p[120:])

短っっ。実行結果は動的計画で計算したものと一致していました。

ついでに母関数の微分を使って平均、分散を出してみます。

from numpy import poly1d as poly

dice = range(7)
total = sum(dice)
prob = [1. * x / total for x in dice]

f = poly(prob[::-1])
f = f ** 30

mean = f.deriv()(1)
sigma = f.deriv()(1) + f.deriv(2)(1) - mean ** 2

print mean
print sigma

因みに今回の問題の場合は各試行が独立なので、平均・分散は以下のように普通に計算できます。結果は母関数から導出したものと同じになりました。

dice = range(1, 7)
tot = sum(dice)

prob = [1. * x / tot for x in dice]

mean = sum([(x + 1) * p for x, p in enumerate(prob)])
sigma = sum([(x + 1) * (x + 1) * p for x, p in enumerate(prob)]) - mean**2

print mean * 30
print sigma * 30