Atcoder Grand Contest 013 E. Placing Squares

Atcoder Grand Contest 013 E. Placing Squares

E: Placing Squares - AtCoder Grand Contest 013 | AtCoder

感想

$\ O(n)\ $を高速化したかったけどできてない.....
想定解はまた今度考える。

解法

まず愚直に$\ DP\ $を書くと$\ O(n^2)\ $になる。この記事ではうまく高速化しないと通せない$\ O(n)\ $解法を説明する。印の処理は容易なので、省略する。

まず$\ dp_i = \sum_{k=0}^{i-1} dp_k * (i - k) ^ 2\ $という式を立てることができるのでこれを変形していく。

$\ dp_i = i^2 * \sum_{k=0}^{i-1} dp_k - 2i * \sum_{k=0}^{i-1} k * dp_k + \sum_{k=0}^{i-1} k^2 * dp_k\ $

$\ dp_i = i^2 * sum_{i-1} - 2 * i * sum2_{i-1} + sum3_{i-1}\ $とおくと

$\ sum_{i} = sum_{i-1} + dp_{i-1}\ $
$\ sum2_{i} = sum2_{i-1} + (i-1)*dp_{i-1}\ $
$\ sum3_{i} = sum3_{i-1} + (i-1)^2 * dp_{i-1}\ $

を用いて順番に計算していける。

ごちゃごちゃ書いたが、要するに総和を知りたいときに一つ前に使った総和そのものが使える典型パターンである。

よってこれは$\ O(n)\ $で計算できる。

さらにこれらは一つ前の情報しか必要としないので、すべて一つの変数を使いまわすことで空間計算量はマークの分の$\ O(n)\ $になる。

実装(5倍くらいTLE(>_<))

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

const int MOD = 1000000007;

int main() {
        int n, m, x;
        scanf(&quot;%d%d&quot;, &n, &m);
        vector<bool> marked(n + 1, false);
        for (int i = 0; i < m; i ++) { scanf(&quot;%d&quot;, &x); marked[x] = true; }
        long long ans = 1;
        long long sum = 1, sum2 = 0, sum3 = 0;
        long long j = 0;
        for (long long i = 1; i <= n; i ++) {
                j += i * 2 - 1; //j = i * i
                j %= MOD;
                if (marked[i]) ans = 0;
                else { 
                        ans = j * sum - 2 * i * sum2 + sum3;
                        ans %= MOD;
                        sum  +=     ans;
                        sum2 += i * ans;
                        sum3 += j * ans;
                        sum  %= MOD;
                        sum2 %= MOD;
                        sum3 %= MOD;
                }
        }
        printf(&quot;%lld\n&quot;, ((ans % MOD) + MOD) % MOD);
        return 0;
}