https://beta.atcoder.jp/contests/arc071/tasks/arc071_d

問題概要

整数Nが与えられたとき、以下の性質をすべて満たす数列が何通りあるか求めよ。

数列は無限長で、各要素は1〜Nのいずれかの整数である
数列の第N項以降はすべて同じ値である
数列のある要素がaであったとき、続くa個の要素はすべて同じ値である (aと同じである必要はない)

N <= 10⁶

解法

左から数列を埋めていくことを考える。数列の性質を考えると、2以上の数のあとにまた2以上の数が来た場合、以降も全部同じ数字で埋まることが確定する(ex. 2, 3 のあとは 3, 3, 3, ... しかありえない)。これを踏まえると以下のDPができる。

dp[i][最後が1かどうか]: i番目まで見て、最後の要素が1もしくは2以上のときの数列の総数

数列をAと表すことにすると、このDPの遷移は以下のようになる。

A[i]が1のとき: 次に何が来てもよい（その決定によって連鎖的にi+2以降の値も決まってしまう、ということがない）。
A[i]が2〜Nで、次が1のとき: 1がA[i]個連続して埋まる。
A[i]が2〜Nで、次が2〜Nのとき: 以降無限に同じ数字が繰り返される。

1, 3の遷移は比較的単純なDPである。2のケースは若干面倒だが要は i -> i+2, i+3, i+4, ..., i+N にそれぞれ同じ値を配ればいいので累積和とかでなんとかなる。

感想

自力でできたので嬉しかった。上の文章は細かい数字のやりくりとか端折ってるんで公式解説とか見たほうがいいと思います。

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, std.bitmanip;

immutable long MOD = 10^^9 + 7;

void main() {
    auto N = readln.chomp.to!int;

    auto dp = new long[][](N+1, 2);
    dp[0][1] = 1;

    auto imos = new long[](3*N);
    long INV = powmod(N-1, MOD-2, MOD);
    long ans = 0;

    foreach (i; 0..N) {
        (dp[i+1][0] += dp[i][1] * (N - 1) % MOD) %= MOD; // 1のあとに2〜N
        (dp[i+1][1] += dp[i][1]) %= MOD; // 1のあとに1
        (dp[N][0] += dp[i][0] * (N - 1) % MOD) %= MOD; // 2〜Nのあとに2〜N (無限ループ)

        // 2〜Nのあとに1
        (imos[i+2] += dp[i][0] * INV % MOD) %= MOD;
        (imos[i+N+1] -= dp[i][0] * INV % MOD) %= MOD;

        (imos[i+1] += imos[i]) %= MOD;
        (dp[i+1][1] += imos[i+1]) %= MOD;
    }

    foreach (i; N..3*N-1) {
        (imos[i+1] += imos[i]) %= MOD;
        (ans += imos[i+1]) %= MOD;
    }

    ans = (ans + dp[N].sum % MOD + MOD) % MOD;
    ans.writeln;
}

long powmod(long a, long x, long m) {
    long ret = 1;
    while (x) {
        if (x % 2) ret = ret * a % m;
        a = a * a % m;
        x /= 2;
    }
    return ret;
}

nebukuro09の日記

AtCoder Regular Contest 071: F - Infinite Sequence

問題概要

解法

感想

コード (D言語)