けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces 552 DIV3 E. Two Teams (R1800)

超苦手タイプ

問題へのリンク

問題概要

 N 要素からなる順列があたえられる。先手と後手が交互に

  • 残っている中から最大要素を選び、その左右  K 要素を合わせて  2K+1 要素を、自分のところに抜き取ってくる (左右に  K 要素残っていない場合はある分だけ抜き取ってくる)

という操作を繰り返す。各要素について、先手と後手のどちらのチームに加わったかを求めよ。

制約

  •  1 \le K \le N \le 2 × 10^{5}

考えたこと

こういうデータ構造色の強い問題はひたすら慣れていくしかないね。

  • 順列の逆順列 (値から添字をとる)
  • 残ってる要素を管理する set
  • 残ってる要素の中で左右がどれかを管理する配列

を持たせることにした。計算時間は  O(N\log{N})

#include <iostream>
#include <vector>
#include <set>
using namespace std;

int main() {
    // 順列と逆順列
    int N, K; cin >> N >> K;
    vector<int> a(N), ia(N);
    for (int i = 0; i < N; ++i) cin >> a[i], --a[i], ia[a[i]] = i;

    // 残っている要素の前と後ろ
    vector<int> prev(N, -1), next(N, -1);
    for (int i = 0; i < N; ++i) prev[i] = i-1, next[i] = i+1;

    // 残っている要素の集合
    set<int> se;
    for (int i = 0; i < N; ++i) se.insert(-i);

    // シミュレーション開始
    int turn = 1;
    vector<int> res(N);
    while (!se.empty()) {
        int ma = -(*se.begin()); se.erase(-ma);
        int ima = ia[ma];
        res[ima] = turn;

        int left = prev[ima], right = next[ima];
        for (int i = 0; i < K; ++i, left = prev[left]) {
            if (left == -1) break;
            int val = a[left];
            se.erase(-val);
            res[left] = turn;
        }
        for (int i = 0; i < K; ++i, right = next[right]) {
            if (right == N) break;
            int val = a[right];
            se.erase(-val);
            res[right] = turn;
        }

        // つなぎ変え
        if (left >= 0) next[left] = right;
        if (right < N) prev[right] = left;

        // 交代
        turn = 3 - turn;
    }
    for (int i = 0; i < N; ++i) cout << res[i];
    cout << endl;
}