mathkyoproの日記

数学や競プロの問題を解説したりします。

部分和の数え上げ

2024/04/10 (3) を追記しました。



次の問題を考えてみましょう。

$1$ 以上 $5N$ 以下の整数の集合 $ \{1, 2, \cdots, 5N\} $ から$0$個以上の整数を選ぶ方法は$2^{5N}$ 通りある。このうち、以下の条件を満たす方法は何通りあるか。

条件: 選んだ整数の集合の要素の和は $5$ で割り切れる。(ただし空集合の要素の和は $0$ と定める)


(1) 競プロの問題として

この問題は、競プロの問題であれば、例えば制約は $1 \le N \le 10000$ が考えられ、また答えが大きくなるので、$998244353$ で割った余りを答えることになるでしょう。この問題は簡単な dp で答えることができます。例えば C++ では AC library を使うと以下のように解をもとめることができます。

#include "bits/stdc++.h"
#include <atcoder/all>
using namespace atcoder;
using namespace std;


signed main() {
	int N; cin >> N;

	//dp[i][j]; i 以下の集合の中から、5で割った余りが j になるように選ぶ選び方。
	vector<vector<modint998244353>> dp(5 * N + 1, vector<modint998244353>(5, 0));
	dp[0][0] = 1; //空集合

	for (int i = 1; i <= 5 * N; ++i) {
		for (int j = 0; j < 5; ++j) {
			dp[i][j] += dp[i - 1][j];
			dp[i][(j + i) % 5] += dp[i - 1][j];
		}
	}

	cout << dp[5 * N][0].val() << endl;
}




(2) 数学の問題として: 形式的冪級数 *1

さて、それでは数学の問題として答えを求められた場合はどうすればよいでしょうか。手で dp のように数え上げるのは難しそうです。dp の考え方を応用するためには、やや天下り的ですが次のような条件を満たす多項式を考えます。

・$x^k$ の係数が、部分集合の和が $k$ になるような選び方となる。
※ 上の dp では計算量削減のため和を $5$ で割った余りを考えていますが、ここではあらわに余りを考えます。

$ \{1, 2, \cdots, k\} $ からの選び方で上記条件を満たすような多項式を $f_{k}(x)$ と置くと、遷移は、$f_{k+1}(x) = (x^{k + 1}+1)f_{k}(x)$ です。$f_{1}(x) = x+1$なので、
$$f_{5N}(x) = (x^{1}+1)(x^{2}+1)\cdots(x^{5N}+1)$$ となります。これは、形式的冪級数とよばれる方法です。求めるべき値は、$f_{5N}(x)$ の $x^{k}$ の係数を $c_{k}$ と置くと、以下のように書けます。$$S_{0} \equiv \sum_{k=0}^{N} c_{5k}$$ また、今後のために、$$S_{r} \equiv \sum_{k=0}^{N} c_{5k + r} \ (r = 0, 1, 2, 3, 4)$$ と定めます。

さて、
\begin{split}
f_{5N}(x) &= (x^{1}+1)(x^{2}+1)\cdots(x^{5N}+1) \\
&= c_{0}x^{0}+c_{1}x^{1}+ \cdots +c_{5N}x^{5N}
\end{split}
から、$c_{5k}$ の値の情報を抜き出すために、$1$ の $5$ 乗根を代入することを考えます。$1$ の原始 $5$ 乗根を $$\zeta = \exp \left( \mathrm{i} \frac{2 \pi}{5} \right) = \cos\left(\frac{2 \pi}{5}\right) + \mathrm{i}\sin\left(\frac{2 \pi}{5}\right)$$ と置くと、$1$ の $5$ 乗根は、$\zeta^{0} = 1, \ \zeta^{1}, \ \zeta^{2}, \ \zeta^{3}, \ \zeta^{4} $ です。
\begin{equation} f_{5N} (\zeta^{0}) = c_{0} \zeta^{0} + c_{1} \zeta^{0} + c_{2} \zeta^{0} + c_{3} \zeta^{0} + c_{4} \zeta^{0} + c_{5} \zeta^{0} + \cdots + c_{5N} \zeta^{0} \tag{1} \end{equation} \begin{equation} f_{5N} (\zeta^{1}) = c_{0} \zeta^{0} + c_{1} \zeta^{1} + c_{2} \zeta^{2} + c_{3} \zeta^{3} + c_{4} \zeta^{4} + c_{5} \zeta^{0} + \cdots + c_{5N} \zeta^{0} \tag{2} \end{equation} \begin{equation} f_{5N} (\zeta^{2}) = c_{0} \zeta^{0} + c_{1} \zeta^{2} + c_{2} \zeta^{4} + c_{3} \zeta^{1} + c_{4} \zeta^{3} + c_{5} \zeta^{0} + \cdots + c_{5N} \zeta^{0} \tag{3} \end{equation} \begin{equation} f_{5N} (\zeta^{3}) = c_{0} \zeta^{0} + c_{1} \zeta^{3} + c_{2} \zeta^{1} + c_{3} \zeta^{4} + c_{4} \zeta^{2} + c_{5} \zeta^{0} + \cdots + c_{5N} \zeta^{0} \tag{4} \end{equation} \begin{equation} f_{5N} (\zeta^{4}) = c_{0} \zeta^{0} + c_{1} \zeta^{4} + c_{2} \zeta^{3} + c_{3} \zeta^{2} + c_{4} \zeta^{1} + c_{5} \zeta^{0} + \cdots + c_{5N} \zeta^{0} \tag{5} \end{equation}
$1$ の $5$ 乗根を代入しているので、$c_{5k}$ にかかるのは、常に $1$ です。他の $c_{5k + 1}, \ c_{5k + 2}, \ c_{5k + 3}, \ c_{5k + 4}$ には、$\zeta^{0} = 1, \ \zeta^{1}, \ \zeta^{2}, \ \zeta^{3}, \ \zeta^{4} $ が交代でかかっていることが分かります。今、$\zeta^{0} = 1, \ \zeta^{1}, \ \zeta^{2}, \ \zeta^{3}, \ \zeta^{4} $ が $1$ の $5$ 乗根であることから、$$ z^5 - 1 = (z-\zeta^{0})(z-\zeta^{1})(z-\zeta^{2})(z-\zeta^{3})(z-\zeta^{4}) \tag{6}$$ と因数分解でき、解と係数の関係から $$\zeta^{0}+\zeta^{1}+\zeta^{2}+\zeta^{3}+\zeta^{4}=0 \tag{7}$$ です。因みに $(7)$ は、$\zeta^{0} = 1, \ \zeta^{1}, \ \zeta^{2}, \ \zeta^{3}, \ \zeta^{4} $ を頂点とする正五角形が、ガウス平面上で 原点を重心とすることと対応しています。以上のことから、$(1) - (5)$ を足すと、$c_{5k}$ 以外は打ち消して、$$ f_{5N} (\zeta^{0}) + f_{5N} (\zeta^{1}) + f_{5N} (\zeta^{2}) + f_{5N} (\zeta^{3}) + f_{5N} (\zeta^{4}) = 5S_{0} \tag{8} $$ を得ます。あとは左辺を求めればよいです。

まず、$f_{5N}(x) = (x^{1}+1)(x^{2}+1)\cdots(x^{5N}+1)$ であることを思い出すと、$$f_{5N}(\zeta^{0}) =2^{5N} \tag{9}$$ です。これは、全ての選び方が $2^{5N}$ 通りであるという自明な事実と対応しています。次に
$$f_{5N}(\zeta^{1}) = \left( (\zeta^{1} + 1)(\zeta^{2} + 1)(\zeta^{3} + 1)(\zeta^{4} + 1)(\zeta^{5} + 1) \right)^{N} \tag{10}$$ ですが、右辺はいくつなのでしょうか。ここで実は $(6)$ に $z=-1$ を代入すると、$(6)$ 右辺は、$(10)$ 右辺の括弧の中の $-1$ 倍であることに気づくと、
$$f_{5N}(\zeta^{1}) = 2^{N}$$ と分かります。同様にして、$$f_{5N}(\zeta^{1}) = f_{5N}(\zeta^{2}) = f_{5N}(\zeta^{3}) = f_{5N}(\zeta^{4}) = 2^{N} \tag{11}$$ を得るので、$(8), \ (9), \ (11)$ から、答えは $$S_{0} = \frac{2^{5N} + 4 \times 2^{N}}{5} \tag{12}$$ と分かります。


ちなみに、他の余りになるような選び方ですが、$ \zeta^{0r} \cdot (1) + \zeta^{1r} \cdot (2) + \zeta^{2r} \cdot (3) + \zeta^{3r} \cdot (4) + \zeta^{4r} \cdot (5)$ を考えることで、$c_{5k + r}$ 以外は打ち消して、$$ \zeta^{0r} \cdot f_{5N} (\zeta^{0}) + \zeta^{1r} \cdot f_{5N} (\zeta^{1}) + \zeta^{2r} \cdot f_{5N} (\zeta^{2}) + \zeta^{3r} \cdot f_{5N} (\zeta^{3}) + \zeta^{4r} \cdot f_{5N} (\zeta^{4}) = 5S_{r}$$ です。ここで、$(9), \ (11)$ から、
$$2^{5N} + (\zeta^{1r} + \zeta^{2r} + \zeta^{3r} + \zeta^{4r}) \times 2^{N} = 5S_{r}$$ ここで、$(7)$ から $$ \zeta^{1r} + \zeta^{2r} + \zeta^{3r} + \zeta^{4r} = - \zeta^{0} = -1 \ (r = 1, 2, 3, 4) $$ であるので、$$S_{1} = S_{2} = S_{3} = S_{4} =\frac{2^{5N} - 2^{N}}{5} \tag{13}$$ となります。


(もう少し dp に寄せた方法)

最初のいくつかは手で dp のように数え上げることがでます。$\{1, 2, 3, 4, 5\}$ の集合から選ぶ方法を数え上げて、$ \mathrm{dp}[5][0] = 8, \ \mathrm{dp}[5][1] = \mathrm{dp}[5][2] = \mathrm{dp}[5][3] = \mathrm{dp}[5][4] = 6$ を得ます。$6$ 以降は、同じことの繰り返しなので、$$g_{5N}(x) = \left( 2(4+3x^{1}+3x^{2}+3x^{3}+3x^{4}) \right)^{N} \tag{14}$$ で $f_{5N}(x)$ と同じように係数を考えればよいです。


※ $f_{5N}(x)$ と $g_{5N}(x)$ は多項式として異なりますが、その $5$ つおきの和が等しいです。言い換えれば、$g_{5N}(x)$ の $x^{k}$ の係数を $d_{k}$とおけば、$$S_{r} \equiv \sum_{k=0}^{N} c_{5k + r} = \sum_{k=0}^{N} d_{5k + r} \ (r = 0, 1, 2, 3, 4)$$ が成り立ちます。


今、$(14)$ 右辺の括弧の中で $x^{1}, \ x^{2}, \ x^{3}, \ x^{4}$ の係数が等しいので、$S_{1} = S_{2} = S_{3} = S_{4}$ と分かります。 より厳密には数学的帰納法により証明できます。$(14)$ に $1, \ \zeta$ を代入して、$$32^{N} = 2^{5N} = S_{0} + S_{1} + S_{2} + S_{3} + S_{4} = S_{0} + 4S_{1} \tag{15}$$ \begin{split}
\left( 2 \left(4+3(\zeta^{1} + \zeta^{2} + \zeta^{3} + \zeta^{4}) \right) \right)^{N} &= S_{0} + \zeta^{1} S_{1} + \zeta^{2} S_{2} + \zeta^{3} S_{3} + \zeta^{4} S_{4} \\
&= S_{0} + (\zeta^{1} + \zeta^{2} + \zeta^{3} + \zeta^{4})S_{1}
\end{split} ここで $(7)$ から、$$2^{N} = S_{0} - S_{1} \tag{16}$$
と分かり、$(15), \ (16)$ から $(12), \ (13)$ を得ます。





(3) 数学の問題として: 行列

さて、(2) で求めた表式を使えば、競プロにおいて例えば繰り返し二乗法を用いて $O \left( \log N \right)$ で答えを求めることができ、$1 \le N \le 2\times10^{17}$ などの制約でも問題を解くことができます。しかし普通競プロにおいて、その制約下でこの問題を解くには、そのようなことはせずに行列を使います *2。数学の問題としても、行列を使って (2) の表式を求めることが出来るのではないでしょうか。
\begin{equation}
\begin{pmatrix} \mathrm{dp}[i][0] \\ \mathrm{dp}[i][1] \\ \mathrm{dp}[i][2] \\ \mathrm{dp}[i][3] \\ \mathrm{dp}[i][4] \end{pmatrix} = A_{i} \begin{pmatrix} \mathrm{dp}[i - 1][0] \\ \mathrm{dp}[i - 1][1] \\ \mathrm{dp}[i - 1][2] \\ \mathrm{dp}[i - 1][3] \\ \mathrm{dp}[i - 1][4] \end{pmatrix} \ \ (i \ge 1)
\end{equation} とすると、
\begin{split}
A_{5}A_{4}A_{3}A_{2}A_{1} &= A_{5}A_{4}A_{3} \begin{pmatrix} 1 &0 &0 &1 &0 \\ 0 &1 &0 &0 &1 \\ 1 &0 &1 &0 &0 \\ 0 &1 &0 &1 &0 \\ 0 &0 &1 &0 &1 \end{pmatrix} \begin{pmatrix} 1 &0 &0 &0 &1 \\ 1 &1 &0 &0 &0 \\ 0 &1 &1 &0 &0 \\ 0 &0 &1 &1 &0 \\ 0 &0 &0 &1 &1 \end{pmatrix} \\
&= A_{5}A_{4}A_{3} \begin{pmatrix} 1 &0 &1 &1 &1 \\ 1 &1 &0 &1 &1 \\ 1 &1 &1 &0 &1 \\ 1 &1 &1 &1 &0 \\ 0 &1 &1 &1 &1 \end{pmatrix} \\
&= A_{5}A_{4} \begin{pmatrix} 1 &0 &1 &0 &0 \\ 0 &1 &0 &1 &0 \\ 0 &0 &1 &0 &1 \\ 1 &0 &0 &1 &0 \\ 0 &1 &0 &0 &1 \end{pmatrix} \begin{pmatrix} 1 &0 &1 &1 &1 \\ 1 &1 &0 &1 &1 \\ 1 &1 &1 &0 &1 \\ 1 &1 &1 &1 &0 \\ 0 &1 &1 &1 &1
\end{pmatrix} \\
&= A_{5}A_{4} \begin{pmatrix} 2 &1 &2 &1 &2 \\ 2 &2 &1 &2 &1 \\ 1 &2 &2 &1 &2 \\ 2 &1 &2 &2 &1 \\ 1 &2 &1 &2 &2 \end{pmatrix} \\
&= A_{5} \begin{pmatrix} 1 &1 &0 &0 &0 \\ 0 &1 &1 &0 &0 \\ 0 &0 &1 &1 &0 \\ 0 &0 &0 &1 &1 \\ 1 &0 &0 &0 &1 \end{pmatrix} \begin{pmatrix} 2 &1 &2 &1 &2 \\ 2 &2 &1 &2 &1 \\ 1 &2 &2 &1 &2 \\ 2 &1 &2 &2 &1 \\ 1 &2 &1 &2 &2 \end{pmatrix} \\
&= A_{5} \begin{pmatrix} 4 &3 &3 &3 &3 \\ 3 &4 &3 &3 &3 \\ 3 &3 &4 &3 &3 \\ 3 &3 &3 &4 &3 \\ 3 &3 &3 &3 &4 \end{pmatrix} \\
&= 2 \begin{pmatrix} 4 &3 &3 &3 &3 \\ 3 &4 &3 &3 &3 \\ 3 &3 &4 &3 &3 \\ 3 &3 &3 &4 &3 \\ 3 &3 &3 &3 &4 \end{pmatrix} \ \ \left( \because A_{5} = 2E \right) \\
\end{split} となります。従って、
\begin{equation}
\begin{pmatrix} \mathrm{dp}[5N][0] \\ \mathrm{dp}[5N][1] \\ \mathrm{dp}[5N][2] \\ \mathrm{dp}[5N][3] \\ \mathrm{dp}[5N][4] \end{pmatrix} = \begin{pmatrix} 8 &6 &6 &6 &6 \\ 6 &8 &6 &6 &6 \\ 6 &6 &8 &6 &6 \\ 6 &6 &6 &8 &6 \\ 6 &6 &6 &6 &8 \end{pmatrix} \begin{pmatrix} \mathrm{dp}[5N - 5][0] \\ \mathrm{dp}[5N - 5][1] \\ \mathrm{dp}[5N - 5][2] \\ \mathrm{dp}[5N - 5][3] \\ \mathrm{dp}[5N - 5][4] \end{pmatrix} \ \ (N \ge 1) \tag{17}
\end{equation} と分かります。これは、式 $(14)$ に対応しています。
\begin{equation}
B = \begin{pmatrix} 4 &3 &3 &3 &3 \\ 3 &4 &3 &3 &3 \\ 3 &3 &4 &3 &3 \\ 3 &3 &3 &4 &3 \\ 3 &3 &3 &3 &4 \end{pmatrix}
\end{equation} とおきます。これのべき乗を求めたいので、この対角化を目指します。$5 \times 5$ の対角化なので難しく感じますが、
\begin{equation}
\begin{pmatrix} 4 - \lambda &3 &3 &3 &3 \\ 3 &4 - \lambda &3 &3 &3 \\ 3 &3 &4 - \lambda &3 &3 \\ 3 &3 &3 &4 - \lambda &3 \\ 3 &3 &3 &3 &4 - \lambda \end{pmatrix} \boldsymbol{v} = \mathbf{0}
\end{equation} なる線型独立なベクトル $\boldsymbol{v}$ を $5$ 本見つければ良いです。これはぐっと眺めると見つけることができます。まず、$\lambda = 16$ に対して、
\begin{equation}
\boldsymbol{v}_{1} = \begin{pmatrix} 1 \\ 1 \\ 1 \\ 1 \\ 1 \end{pmatrix}
\end{equation} が条件を満たします。また、$\lambda = 1$ に対して、
\begin{equation}
\boldsymbol{v}_{2} = \begin{pmatrix} 1 \\ -1 \\ 0 \\ 0 \\ 0 \end{pmatrix}, \ \boldsymbol{v}_{3} = \begin{pmatrix} 1 \\ 0 \\ -1 \\ 0 \\ 0 \end{pmatrix}, \ \boldsymbol{v}_{4} = \begin{pmatrix} 1 \\ 0 \\ 0 \\ -1 \\ 0 \end{pmatrix}, \ \boldsymbol{v}_{4} = \begin{pmatrix} 1 \\ 0 \\ 0 \\ 0 \\ -1 \end{pmatrix}
\end{equation} の 4本が条件を満たします。従って、
\begin{equation}
P = \begin{pmatrix} \boldsymbol{v_1} \ \boldsymbol{v_2} \ \boldsymbol{v_3} \ \boldsymbol{v_4} \ \boldsymbol{v_5}\end{pmatrix} = \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-1 &0 &0 &0 \\ 1 &0 &-1 &0 &0 \\ 1 &0 &0 &-1 &0 \\ 1 &0 &0 &0 &-1 \end{pmatrix} \tag{18}
\end{equation} とおけば、
\begin{equation}
BP = P \begin{pmatrix} 16 &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 \\ 0 &0 &1 &0 &0 \\ 0 &0 &0 &1 &0 \\ 0 &0 &0 &0 &1 \end{pmatrix}
\end{equation} であるので、
\begin{equation}
P^{-1}BP = \begin{pmatrix} 16 &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 \\ 0 &0 &1 &0 &0 \\ 0 &0 &0 &1 &0 \\ 0 &0 &0 &0 &1 \end{pmatrix}
\end{equation} です。従って、
\begin{equation}
(P^{-1}BP)^N = P^{-1}B^{N}P = \begin{pmatrix} 16^N &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 \\ 0 &0 &1 &0 &0 \\ 0 &0 &0 &1 &0 \\ 0 &0 &0 &0 &1 \end{pmatrix}
\end{equation} となります。以上より、
\begin{equation}
B^{N} = P \begin{pmatrix} 16^N &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 \\ 0 &0 &1 &0 &0 \\ 0 &0 &0 &1 &0 \\ 0 &0 &0 &0 &1 \end{pmatrix} P^{-1} \tag{19}
\end{equation} を得ます。これを求めるには、$P^{-1}$ を求めなければなりません。$5 \times 5$ の逆行列なので難しく感じますが、
\begin{equation}
PP^{-1} = \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-1 &0 &0 &0 \\ 1 &0 &-1 &0 &0 \\ 1 &0 &0 &-1 &0 \\ 1 &0 &0 &0 &-1 \end{pmatrix} P^{-1} = E
\end{equation} を満たす行列 $P^{-1}$ は、ぐっと眺めると $1$ 行ずつ決めることができます。すると、
\begin{equation}
P^{-1} = \frac{1}{5} \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-4 &1 &1 &1\\ 1 &1 &-4 &1 &1 \\ 1 &1 &1 &-4 &1 \\ 1 &1 &1 &1 &-4 \end{pmatrix} \tag{20}
\end{equation} と分かります。以上より、式 $(18) - (20)$ から、
\begin{split}
B &= \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-1 &0 &0 &0 \\ 1 &0 &-1 &0 &0 \\ 1 &0 &0 &-1 &0 \\ 1 &0 &0 &0 &-1 \end{pmatrix} \begin{pmatrix} 16^N &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 \\ 0 &0 &1 &0 &0 \\ 0 &0 &0 &1 &0 \\ 0 &0 &0 &0 &1 \end{pmatrix} \frac{1}{5} \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-4 &1 &1 &1\\ 1 &1 &-4 &1 &1 \\ 1 &1 &1 &-4 &1 \\ 1 &1 &1 &1 &-4 \end{pmatrix} \\
&= \frac{1}{5} \begin{pmatrix} 16^N &1 &1 &1 &1 \\ 16^N &-1 &0 &0 &0 \\ 16^N &0 &-1 &0 &0 \\ 16^N &0 &0 &-1 &0 \\ 16^N &0 &0 &0 &-1 \end{pmatrix} \begin{pmatrix} 1 &1 &1 &1 &1 \\ 1 &-4 &1 &1 &1\\ 1 &1 &-4 &1 &1 \\ 1 &1 &1 &-4 &1 \\ 1 &1 &1 &1 &-4 \end{pmatrix} \\
&= \frac{1}{5} \begin{pmatrix} a_n &b_n &b_n &b_n &b_n \\ b_n &a_n &b_n &b_n &b_n \\ b_n &b_n &a_n &b_n &b_n \\ b_n &b_n &b_n &a_n &b_n \\ b_n &b_n &b_n &b_n &a_n \end{pmatrix}
\end{split} です。但し、$a_n = 16^N + 4, \ b_n = 16^N - 1$ と置きました。式 $(17)$ から、
\begin{split}
\begin{pmatrix} \mathrm{dp}[5N][0] \\ \mathrm{dp}[5N][1] \\ \mathrm{dp}[5N][2] \\ \mathrm{dp}[5N][3] \\ \mathrm{dp}[5N][4] \end{pmatrix} &= \left( 2B \right)^{N} \begin{pmatrix} \mathrm{dp}[0][0] \\ \mathrm{dp}[0][1] \\ \mathrm{dp}[0][2] \\ \mathrm{dp}[0][3] \\ \mathrm{dp}[0][4] \end{pmatrix} \\
&= \left( 2B \right)^{N} \begin{pmatrix} 1 \\ 0 \\ 0 \\ 0 \\ 0 \end{pmatrix} \\
&= \frac{2^N}{5} \begin{pmatrix} a_n \\ b_n \\ b_n \\ b_n \\ b_n \end{pmatrix} \\
&= \frac{1}{5} \begin{pmatrix} 32^N + 4 \cdot 2^N \\ 32^N - 2^N \\ 32^N - 2^N \\ 32^N - 2^N \\ 32^N - 2^N \end{pmatrix} \\
\end{split} これは、式 $(12), (13)$ を意味しており、解を得ました。

(2) より計算は煩雑ですが、考えるところや工夫は少なくて済みます。ごり押し解法という感じです。