Atcoder Grand Contest 013 E. Placing Squares(追記）

Atcoder Grand Contest 013 E. Placing Squares(追記)

E: Placing Squares - AtCoder Grand Contest 013 | AtCoder

感想

愚直$\ DP\ $の高速化に失敗したので、解説放送を見て想定解を実装した。

解法

$\ 0\ $から$\ n\ $までの任意の位置に柵をおいて、その柵で囲まれたところそれぞれに2色のボールを置く（同じ位置でも良い）ときの場合の数を求めればそれが答えである。

解説放送ではこれに気づけばあとは自明な$\ DP\ $で解けると言っていたが、この遷移で少し悩んだので丁寧に書く。

まず$\ dp[i][j] := \ $（区間$\ [0, i)\ $の配置が完了していて、完了している区間の中で一番右の柵より右側に置いてあるボールの個数が$\ j\ $個となっている場合の数）とする。

ここで注意することは、区間$\ i\ $の柵の状態は無視することである。位置$\ i - 1\ $と$\ i\ $の間の位置$\ P_i\ $に置くボールの個数と、位置$\ i - 1\ $に柵を置くかどうかで場合分けをして遷移を考える。

まず位置$\ i - 1\ $に印がない場合を考える（すなわち柵を置いてもよいし置かなくても良い）。例えば$\ dp[i][0]\ $の値は次のようになる。

$\ i - 1\ $に柵を置く場合、$\ dp[i][0]\ += dp[i - 1][2]\ $（$\ P_i\ $にボールを置かない）
$\ i - 1\ $に柵を置かない場合、$\ dp[i][0]\ += dp[i - 1][0]\ $（$\ P_i\ $にボールを置かない）

つまり、常に上の$\ dp[i][j]\ $の条件に矛盾しないように、ボールを置いて考えればよい。同様にすると、次のようになる。

$\ i - 1\ $に柵を置く場合、$\ dp[i][1]\ += 2 * dp[i - 1][2]\ $（$\ P_i\ $にどちらかのボールを選んで置く）
$\ i - 1\ $に柵を置かない場合、$\ dp[i][1]\ += 2 * dp[i - 1][0]\ $（$\ P_i\ $にどちらかのボールを選んで置く）$+ dp[i - 1][1]\ $（$\ P_i\ $にボールを置かない）

$\ i - 1\ $に柵を置く場合、$\ dp[i][2]\ += dp[i - 1][2]\ $（$\ P_i\ $に両方のボールを置く）
$\ i - 1\ $に柵を置かない場合、$\ dp[i][2]\ += dp[i - 1][0]\ $（$\ P_i\ $に両方のボールを置く）$+ dp[i - 1][1]\ $（すでに置いてあるボールとは異なる色のボールを$\ P_i\ $に置く）$+ dp[i - 1][2]\ $（$\ P_i\ $に置かない）

以上から、次が従う。

$dp[i][0] = 1 * dp[i - 1][0] + 0 * dp[i - 1][1] + 1 * dp[i - 1][2]\ $
$dp[i][1] = 2 * dp[i - 1][0] + 1 * dp[i - 1][1] + 2 * dp[i - 1][2]\ $
$dp[i][2] = 1 * dp[i - 1][0] + 1 * dp[i - 1][1] + 2 * dp[i - 1][2]\ $

これは明らかに次のように書ける。

$
\left(
\begin{array}{rrr}
dp[i][0] \\
dp[i][1] \\
dp[i][2] \\
\end{array}
\right) = \left(
\begin{array}{rrr}
1 & 0 & 1 \\
2 & 1 & 2 \\
1 & 1 & 2 \\
\end{array}
\right) \left(
\begin{array}{rrr}
dp[i - 1][0] \\
dp[i - 1][1] \\
dp[i - 1][2] \\
\end{array}
\right)
$

次に位置$\ i - 1\ $に印がついている場合を考えると、その位置では柵が置けないため、上の場合分けにおける柵を置かない場合のみを考えることで

$dp[i][0] = 1 * dp[i - 1][0] + 0 * dp[i - 1][1] + 0 * dp[i - 1][2]\ $
$dp[i][1] = 2 * dp[i - 1][0] + 1 * dp[i - 1][1] + 0 * dp[i - 1][2]\ $
$dp[i][2] = 1 * dp[i - 1][0] + 1 * dp[i - 1][1] + 1 * dp[i - 1][2]\ $

であることがわかるため、以下のようになる。

$
\left(
\begin{array}{rrr}
dp[i][0] \\
dp[i][1] \\
dp[i][2] \\
\end{array}
\right) = \left(
\begin{array}{rrr}
1 & 0 & 0 \\
2 & 1 & 0 \\
1 & 1 & 1 \\
\end{array}
\right) \left(
\begin{array}{rrr}
dp[i - 1][0] \\
dp[i - 1][1] \\
dp[i - 1][2] \\
\end{array}
\right)
$

よって、各印の間は上の行列の累乗を用いて計算し、印の位置で下の行列を1回かけるということを繰り返せば高速に答えを求めることができる。

実装

#include "bits/stdc++.h"
using namespace std;

const int MOD = 1000000007;

typedef long long Type;
typedef vector<vector<Type>> Matrix;
Matrix Mul(const Matrix &a, const Matrix &b) {
        assert(a[0].size() == b.size());
        Matrix c(a.size(), vector<Type> (b[0].size()));
        for (int i = 0; i < a.size(); i ++) {
                for (int k = 0; k < b.size(); k ++) {
                        for (int j = 0; j < b[0].size(); j ++) {
                                c[i][j] = (c[i][j] + a[i][k] * b[k][j]) % MOD;
                        }
                }
        }
        return c;
}
Matrix Pow(Matrix a, long long n) {
        assert(a.size() == a[0].size());
        Matrix b(a.size(), vector<Type> (a.size()));
        for (int i = 0; i < a.size(); i ++) {
                b[i][i] = 1;
        }
        while (n > 0) {
                if (n & 1) b = Mul(b, a);
                a = Mul(a, a);
                n >>= 1;
        }
        return b;
}

int main() {
        int n, m, x;
        scanf("%d%d", &n, &m);
        Matrix a = {{1, 0, 0}, 
                    {2, 1, 0}, 
                    {1, 1, 1}};
        Matrix b = {{1, 0, 1},
                    {2, 1, 2},
                    {1, 1, 2}};
        Matrix ans = a;
        int prev = 1;
        while (m --) {
                scanf("%d", &x);
                ans = Mul(Pow(b, x - prev), ans);
                ans = Mul(a, ans);
                prev = x + 1;
        }
        ans = Mul(Pow(b, n - prev), ans);
        printf("%lld\n", ans[2][0]);
        return 0;
}