問題
整数 $n$ が与えられる。$a _ i > b _ i \geq 0$, $\sum _ {i = 1} ^ {k} (a _ i + b _ i) \leq n$ をみたす $k$ 組の整数のペア $(a _ i, b _ i)$ すべてについて $\prod _ {i = 1} ^ {k} (a _ i - b _ i)$ の総和を求めよ。($n \leq 10 ^ 9, k = 5$)
解説
$c _ i = a _ i - b _ i$ とおくと、$b _ i \geq 0$, $c _ i > 0$, $\sum _ {i = 1} ^ {k} (2 b _ i + c _ i) \leq n$ をみたす $k$ 組の整数のペア $(b _ i, c _ i)$ すべてについて $\prod _ {i = 1} ^ {k} c _ i$ の総和を求めよ、と言い換えられる。このとき、まず $b _ i \geq 0$ を選んでから $x = \sum _ {i = 1} ^ {k} b _ i$ とおいて、$c _ i > 0$, $\sum _ {i = 1} ^ {k} c _ i \leq n - 2 x$ をみたす $c _ i$ を $b _ i$ に関係なく選んでよいので、各 $x$ に対する $b _ i$ の組合せの数を $f(x)$, $\prod _ {i = 1} ^ {k} c _ i$ の総和を $g(x)$ とおき、すべての $x$ に対して $f(x) \times g(x)$ の総和を求めればよい。
以下では、二項係数 ${}_{a} \mathrm{C} _ {b}$ は $a<b$ の場合には $0$ と定義する。
まず、$b _ i \geq 0$, $\sum _ {i = 1} ^ {k} b _ i = x$ をみたす $b _ i$ の選び方は $x$ 個のボールと $k - 1$ 個の仕切りを並べる順列の総数に等しいから、 $$ f(x) = {} _ {x + k -1} \mathrm{C} _ {k - 1} $$
次に、$c _ i > 0$, $\sum _ {i = 1} ^ {k} c _ i \leq n - 2 x$ をみたす $c _ i$ について $\prod _ {i = 1} ^ {k} c _ i$ の総和を求める。$c _ i \geq 0$ と考えても積には影響がないので、一般に、「$c _ i \geq 0$, $\sum _ {i = 1} ^ {k} c _ i \leq y$ をみたす $c _ i$ について $\prod _ {i = 1} ^ {k} c _ i$ の総和を求める」という問題について考える。
$y$ 個のボールと $k$ 個の仕切りを並べた順列と $c _ i$ の選び方は1対1に対応している。このとき $k + 1$ 個のブロックができるが、最後の $1$ 個を無視した $k$ 個のブロックから $1$ 個ずつボールを選ぶ組合せの数と $\prod _ {i = 1} ^ {k} c _ i$ が等しい。よって、$\prod _ {i = 1} ^ {k} c _ i$ の総和は $y + k$ 個から $2 k$ 個選ぶ組合せの数に等しく、 $$ g(x) = {} _ {n - 2 x + k} \mathrm{C} _ {2 k} $$
$f(x),g(x)$ の次数はそれぞれ $k - 1, 2k$ なので、$P(n) = \sum _ {x = 0} ^ {n} f(x) g(x)$ の次数は $3 k$ となる。ただし、$P(n)$ は $n$ の偶奇によって異なる多項式となることに注意する。例えば、$k = 1$ のとき $$ P(n) = \sum _ {x = 0} ^ {n} {}_{n - 2 x + 1} \mathrm{C} _ {2} $$ であり、$n$ が偶数なら $$ P(n) = {} _ {n + 1} \mathrm{C} _ {2} + {} _ {n - 1} \mathrm{C} _ {2} + \cdots + {} _ {3} \mathrm{C} _ {2} $$ だが、$n$ が奇数なら $$ P(n) = {} _ {n + 1} \mathrm{C} _ {2} + {} _ {n - 1} \mathrm{C} _ {2} + \cdots + {} _ {2} \mathrm{C} _ {2} $$
一般に、$d$ 次の多項式 $F(x)$ について、異なる $d + 1$ 個の $x _ i$ が与えられたとき、 $$ F(x) = \sum _ {i = 1} ^ {d + 1} F(x _ i) \frac{\prod _ {j \neq i} (x - x _ {j})}{\prod _ {j \neq i} (x _ {i} - x _ {j})} $$ が成り立つ。(Lagrangeの補間公式)
以上より、一般の $n$ に対して $P(n)$ を求めるには、$n = 0, 2, \ldots, 6 k$ と $n = 1, 3, \ldots, 6 k + 1$ のそれぞれについて $P(n)$ を計算しておいて上の公式を使えばよい。
実装
Lagrange補間
lst に $(x _ i, F(x _ i))$ のリストを渡して、lag = Lagrange(lst) というオブジェクトを作り、 lag.pln(x) で補間された $F(x)$ を返す。
剰余環上での計算を考えているので、割り算の部分は素数modによる逆元をかけている。
class Lagrange: def __init__(self, lst): self.lst = lst def prd(self, j, x): tmp = 1 for i, (xi, yi) in enumerate(self.lst): if j != i: tmp *= (x - xi) tmp %= mod return tmp def pln(self, x): tmp = 0 for i, (xi, yi) in enumerate(self.lst): tmp += yi * (self.prd(i, x) * pow(self.prd(i, xi), mod - 2, mod)) tmp %= mod return tmp
二項係数と$P(n)$ の定義
def cmb(x, y): if x < y: return 0 tmp = 1 for i in range(y): tmp *= x - i tmp *= pow(y - i, mod - 2, mod) tmp %= mod return tmp def pln(n): tmp = 0 for x in range(n): tmp += cmb(x + k - 1, k - 1) * cmb(n - (2 * x) + k, 2 * k) tmp %= mod return tmp
計算
偶奇に分けてLagrange補間を適用し、各入力 $n$ に対して $P(n)$ を求めている。
mod = 10**9 + 7 k = 5 lst_eve = [(i, pln(i)) for i in range(0, 6 * k + 1, 2)] lst_odd = [(i, pln(i)) for i in range(1, 6 * k + 2, 2)] lag_eve = Lagrange(lst_eve) lag_odd = Lagrange(lst_odd) t = int(input()) for _ in range(t): n = int(input()) if n % 2 == 0: print(lag_eve.pln(n)) else: print(lag_odd.pln(n))
