Codeforces Round #641 (Div. 2) D - Orac and Medians 解説

問題へのリンク

問題概要

長さ $N$ の数列 $A$ について、以下の操作を繰り返し全ての要素を $K$ にできるか判定せよ。

数列の区間について、全ての要素をその区間の中央値に置き換える。ただし、中央値は区間の長さを $L$ とすると、小さい方から $\lfloor \frac{L+1}{2} \rfloor$ とする。

解説

色々実験してみると、大きい操作の前に小さい操作で代替の操作が実現可能であることがわかる。

区間の長さを $2$ としたとき、その区間に操作をすると、その区間の要素を全て小さい方に変えることができる。

以下のような数列に置き換えて考える。（中央値であり、 $K$ に合わせられるかを考える問題なので、 $K$ との大小関係以外の情報は無駄である）

$A_i \lt K$ のとき $B_i = 0$
$A_i == K$ のとき $B_i = 1$
$A_i \gt K$ のとき $B_i = 2$

まず、少なくとも1つは $1$ がないといけないので、それを仮定する。

このとき、連続する $2$ 要素が $\left\{ 1,1 \right\}, \left\{ 1,2 \right\}$ のとき、 $\left\{ 1,1 \right\}$ にでき、そこを起点として全ての要素を $1$ にできる。

連続する $3$ 要素が $\left\{ 0,1,1 \right\}, \left\{ 0,1,2 \right\},\left\{ 1,1,1 \right\}, \left\{ 1,1,2 \right\},\left\{ 1,2,2 \right\}$ のときも同様である。

また、すべての $1$ の隣が $0$ の場合、 $1$ の隣を $2$ にすることができればそこを起点として同様に達成可能である。

$1$ の隣に $2$ を持ってくるために必要なのは、連続する2要素が $\left\{ 2,2 \right\}$ 、または連続する3要素が $\left\{ 1,2,2 \right\}, \left\{ 2,2,2 \right\}$ の場合である。それらの区間を起点として、 $2$ を $1$ の隣まで移動させればよい。

長さ $4$ の区間については、上記の条件を確認することと同様なので、 $3$ まで確認すれば十分である。同様に、長さ $5$ 以上についても言えるので、長さ $3$ までの条件を考えればよいことになる。

上記の条件をまとめると、

少なくとも一つは $1$ が存在する
ある連続する3要素が $\left\{ 0,1,1 \right\}, \left\{ 0,1,2 \right\},\left\{ 0,2,2 \right\}, \left\{ 1,1,1 \right\}, \left\{ 1,1,2 \right\},\left\{ 1,2,2 \right\}, \left\{ 2,2,2 \right\}$ である

となるので、それを実装すればよい。

void solve() {
    int N, K;
    cin >> N >> K;
    vec<int> A(N), B(N);
    rep(i, N) {
        cin >> A[i];
        if (A[i] < K) B[i] = 0;
        if (A[i] == K)B[i] = 1;
        if (A[i] > K) B[i] = 2;
    }
    bool exist = false;
    rep(i, N)if (B[i] == 1)exist = true;
    //一個もKがなかったらダメ
    if (not exist) {
        cout << "no" << endl;
        return;
    }
    //1個以上はある状態
    bool ok = false;
    if (N == 1)ok = true;
    else if (N == 2) {
        sort(all(B));
        if (B[0] && B[1]) ok = true;
    }
    // N > 3
    else {
        for (int i = 0; i + 2 < N; i++) {
            vector<int> C;
            C.push_back(B[i]);
            C.push_back(B[i + 1]);
            C.push_back(B[i + 2]);
            sort(all(C));
            if(C[0] == 0 && C[1] == 1 && C[2] == 1)ok = true;
            if(C[0] == 0 && C[1] == 1 && C[2] == 2)ok = true;
            if(C[0] == 0 && C[1] == 2 && C[2] == 2)ok = true;
            if(C[0] == 1 && C[1] == 1 && C[2] == 1)ok = true;
            if(C[0] == 1 && C[1] == 1 && C[2] == 2)ok = true;
            if(C[0] == 1 && C[1] == 2 && C[2] == 2)ok = true;
            if(C[0] == 2 && C[1] == 2 && C[2] == 2)ok = true;
        }
    }

    cout << (ok ? "yes" : "no") << endl;
}

int main() {
    int t;
    cin >> t;
    while (t--) solve();
}