2020-11-01

AtCoder ARC 107 E - Mex Mat (橙色, 800 点)

コンテスト本番はこの問題から解いた！！確信を持つのに時間かかった

問題概要

0 と 1 と 2 のみからなる $N \times N$ の行列 $a$ がある。そのうちの $a_{1,j}$ と $a_{i,1}$ の値のみがわかっている。他のマスの値は、 $i \ge 2, j \ge 2$ に対して

$a_{i, j}$ = mex( $a_{i-1, j}, a_{i, j-1}$ )

と定められている。 $a$ 全体の 0, 1, 2 の個数をそれぞれ求めよ。

制約

$1 \le N \le 5 \times 10^{5}$

考えたこと

この問題なら、単純解法との比較など、いっぱい実験してから提出できそうということで、潜伏のリスクは低いと見た。そして雰囲気得意そうだから確実に通そうと思った。

まずはあれこれ手を動かして実験してみることにした (本当は初手 PC 実験の方がよいのかもしれない)。そうすると、しばらくすると

「ある程度進んだら、右下方向への斜めラインの数値は同じ数値が並ぶ」

という特徴が見えてきた。具体的には、最初の 3 行 3 列分を埋めた後、4 行 4 列目以降については成り立つ模様。そして確かに、任意の $a_{1, 2}, a_{1, 3}, a_{1, 4}, a_{2, 1}, a_{3, 1}, a_{4, 1}$ の設定値に対して

$a_{3, 3} = a_{4, 4}$

が成立することが言える。このことから直ちに、任意の $i \ge 3, j \ge 3$ に対して

$a_{i, j} = a_{i+1, j+1}$

が成立することも言える。このことを利用して解く。 $O(N)$ の計算量で解ける。

コード

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

inline int mex(int i, int j) {
    if (i > j) swap(i, j);
    if (i >= 1 && j >= 1) return 0;
    if (j == 0) return 1;
    else if (j == 1) return 2;
    else return 1;
}

vector<long long> solve(deque<int> a, deque<int> b) {
    vector<long long> res(3, 0);
    int N = a.size();
    for (int iter = 0; iter < min(3, (int)b.size()); ++iter) {
        for (int i = 0; i < a.size(); ++i) res[a[i]]++;
        for (int i = 0; i < b.size(); ++i) res[b[i]]++;
        a[0] = mex(a[1], b[0]);
        for (int i = 1; i+1 < a.size(); ++i) a[i] = mex(a[i-1], a[i+1]);
        a.pop_back();
        b[0] = mex(b[1], a[0]);
        for (int i = 1; i+1 < b.size(); ++i) b[i] = mex(b[i-1], b[i+1]);
        b.pop_back();
    }
    vector<long long> add(3, 0);
    for (int i = 0; i < a.size(); ++i) add[a[i]]++;
    for (int i = 0; i < b.size(); ++i) add[b[i]]++;
    while (!a.empty()) {
        for (int i = 0; i < 3; ++i) res[i] += add[i];
        if (!a.empty()) add[a.back()]--, a.pop_back();
        if (!b.empty()) add[b.back()]--, b.pop_back();
    }
    return res;
}

int main() {
    int N;
    cin >> N;
    deque<int> a(N), b(N-1);
    for (int i = 0; i < N; ++i) cin >> a[i];
    for (int i = 0; i < N-1; ++i) cin >> b[i];

    auto res = solve(a, b);
    cout << res[0] << " " << res[1] << " " << res[2] << endl;
}

コンテスト中のランダムケース生成や単純解法との比較も混みのコード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

#define COUT(x) cout << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl
template<class T1, class T2> ostream& operator << (ostream &s, pair<T1,T2> P)
{ return s << '<' << P.first << ", " << P.second << '>'; }
template<class T> ostream& operator << (ostream &s, vector<T> P)
{ for (int i = 0; i < P.size(); ++i) { if (i > 0) { s << " "; } s << P[i]; } return s; }
template<class T> ostream& operator << (ostream &s, vector<vector<T> > P)
{ for (int i = 0; i < P.size(); ++i) { s << endl << P[i]; } return s << endl; }
template<class T> ostream& operator << (ostream &s, set<T> P)
{ for(auto it : P) { s << "<" << it << "> "; } return s << endl; }
template<class T1, class T2> ostream& operator << (ostream &s, map<T1,T2> P)
{ for(auto it : P) { s << "<" << it.first << "->" << it.second << "> "; } return s << endl; }

unsigned int randInt() {
    static unsigned int tx = 123456789, ty=362436069, tz=521288629, tw=88675123;
    unsigned int tt = (tx^(tx<<11));
    tx = ty; ty = tz; tz = tw;
    return ( tw=(tw^(tw>>19))^(tt^(tt>>8)) );
}

int randInt(int minv, int maxv) {
    return randInt() % (maxv - minv + 1) + minv;
}

inline int mex(int i, int j) {
    if (i > j) swap(i, j);
    if (i >= 1 && j >= 1) return 0;
    if (j == 0) return 1;
    else if (j == 1) return 2;
    else return 1;
}

vector<long long> naive(deque<int> a, deque<int> b) {
    vector<long long> res(3, 0);
    int N = a.size();
    vector<vector<int>> fi(N, vector<int>(N));
    for (int i = 0; i < N; ++i) fi[0][i] = a[i];
    for (int i = 1; i < N; ++i) fi[i][0] = b[i-1];
    for (int i = 1; i < N; ++i) {
        for (int j = 1; j < N; ++j) {
            fi[i][j] = mex(fi[i-1][j], fi[i][j-1]);
        }
    }
    for (int i = 0; i < N; ++i) for (int j = 0; j < N; ++j) res[fi[i][j]]++;
    return res;
}

long long sum(vector<long long> a) {
    long long res = 0;
    for (int i = 0; i < a.size(); ++i) res += a[i];
    return res;
}

vector<long long> solve(deque<int> a, deque<int> b) {
    vector<long long> res(3, 0);
    int N = a.size();
    for (int iter = 0; iter < min(10, (int)b.size()); ++iter) {
        for (int i = 0; i < a.size(); ++i) res[a[i]]++;
        for (int i = 0; i < b.size(); ++i) res[b[i]]++;
        a[0] = mex(a[1], b[0]);
        for (int i = 1; i+1 < a.size(); ++i) a[i] = mex(a[i-1], a[i+1]);
        a.pop_back();
        b[0] = mex(b[1], a[0]);
        for (int i = 1; i+1 < b.size(); ++i) b[i] = mex(b[i-1], b[i+1]);
        b.pop_back();
    }
    vector<long long> add(3, 0);
    for (int i = 0; i < a.size(); ++i) add[a[i]]++;
    for (int i = 0; i < b.size(); ++i) add[b[i]]++;
    while (!a.empty()) {
        for (int i = 0; i < 3; ++i) res[i] += add[i];
        if (!a.empty()) add[a.back()]--, a.pop_back();
        if (!b.empty()) add[b.back()]--, b.pop_back();
    }
    return res;
}

int main() {
    /*
    int N_ = 3000;
    deque<int> a_(N_, 1), b_(N_-1, 1);
    for (int i = 100; i < N_; ++i) a_[i] = randInt(0, 2);
    for (int i = 100; i < N_-1; ++i) b_[i] = randInt(0, 2);
    COUT(naive(a_, b_));
    COUT(solve(a_, b_));
*/

    int N;
    cin >> N;
    deque<int> a(N), b(N-1);
    for (int i = 0; i < N; ++i) cin >> a[i];
    for (int i = 0; i < N-1; ++i) cin >> b[i];

    auto res = solve(a, b);
    cout << res[0] << " " << res[1] << " " << res[2] << endl;
}