AtCoder Grand Contest 044 B - Joker

問題概要

$N \times N$ の正方形の座席を持つ劇場に $N^2$ 人の客が座っている。客が席を立つ順番が与えられるので，以下を求めよ。

客が自席から外周へと出るまでに通る，人の座っている座席の数の，すべての客についての総和。

解法概要

愚直にシミュレートすると時間計算量が $O(N^4)$ となってしまい，間に合わない。しかし，最初の状態における各座席の客が外周に出るまでに通る席の数に着目すると，例えば $N = 6$ なら

である。もう少し正確に言うと，$N$ に対する「最初の状態における各座席の客が外周に出るまでに通る席の数」 (以降，「初期コスト」と表現する) の総和は $N$ が偶数なら $(N-2) ^2 + (N-4)^2 + \cdots + 2^2$ ，奇数なら $(N-2) ^2 + (N-4)^2 + \cdots + 1^2$ である。$\sum_{k=1}^n k^2 = \frac{1}{6} n(n+1)(2n+1)$ であることから，初期コストの総和も $O(n^3)$ で抑えられる。

ここから，ある客が席を立ったときに，その影響範囲の数のみ減らせるような解法であれば通りそうだと考えられる。

ある客が席を立ったときに，影響を受ける範囲でだけ数を減らすには，席を立った客の席から BFS や DFS でコスト (外周に出るまでに通る，人が座っている席の数) が減るような席について値を更新していけば良い。ある席でコストが削減できなかったときに，その先でもコスト削減はできていないはずなのでその先は見に行かなくていい。コストの更新はソースコード参照。まだ人が座っている席を通るたびに値を増やすイメージ。

ソースコード

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>

using namespace std;

constexpr int dyx[4][2] = {
    { 0, 1}, {-1, 0}, {0,-1}, {1, 0}
};

int main() {
    int n;
    cin >> n;
    auto f = [n](int y, int x) { return y * n + x; };
    vector<int> p(n*n), cost(n*n, 0);
    vector<bool> sit(n*n, true);
    for(int i=0;i<n*n;++i) {
        cin >> p[i];
        --p[i];
    }
    for(int i=0;i<n*n;++i) {
        int y = i / n, x = i % n;
        cost[i] = min({y, n-1-y, x, n-1-x});
    }

    int ans = 0LL;
    for(int i=0;i<n*n;++i) {
        ans += cost[p[i]];
        
        queue<int> que;
        que.push(p[i]);
        sit[p[i]] = false;
        while(!que.empty()) {
            int j = que.front(); que.pop();
            int y = j / n, x = j % n;
            for(int k=0;k<4;++k) {
                int yy = y + dyx[k][0], xx = x + dyx[k][1];
                if(0 <= yy && yy < n && 0 <= xx && xx < n && cost[f(yy,xx)] > cost[f(y,x)] + int(sit[f(y,x)])) {
                    cost[f(yy,xx)] = cost[f(y,x)] + int(sit[f(y,x)]);
                    que.push(f(yy,xx));
                }
            }
        }
    }

    cout << ans << endl;
}

所感

解けてしまえば青 diff も納得なのだが，本番はまったくわからなかった。

それにしても，社会人になってから週1で解説記事を書く習慣が崩壊してるな。