AtCoder ABC 380 C - Move Segment (5Q, 灰色, 300 点)

ランレングス圧縮で解いた

• 問題へのリンク

問題概要

0 と 1 からなる長さ の文字列 が与えられる。この文字列の中での「1 の塊」のうち、左から 番目のものを、 番目のものの右に移動させよ。

例："010011100011001" → "010011111000001"

制約

• 

考えたこと

まず、文字列 をランレングス圧縮しよう。そうすると、

"010011100011001"
→ (0, 1), (1, 1), (0, 2), (1, 3), (0, 3), (1, 2), (0, 2), (1, 0)

というようになる。ランレングス圧縮のやり方については、次の問題の公式解説を見るとよい。

atcoder.jp

さて、ランレングス圧縮できれば、ランレングス圧縮した列上で、 番目の「1 の塊」と、その左の「0 の塊」を swap すればよい。上の例では、

(0, 1), (1, 1), (0, 2), (1, 3), (0, 3), (1, 2), (0, 2), (1, 0)
→ (0, 1), (1, 1), (0, 2), (1, 3), (1, 2), (0, 3), (0, 2), (1, 0)

というようになる。こうして得られた列を再び 0-1 文字列に直せば OK。

計算量は となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pint = pair<int,int>;

int main() {
    int N, K;
    string S;
    cin >> N >> K >> S;

    vector<pint> lens;
    for (int i = 0; i < N; ) {
        int j = i;
        while (j < N && S[j] == S[i]) j++;
        lens.emplace_back((int)(S[i]-'0'), j - i);
        i = j;
    }

    int num = 0;
    for (int i = 0; i < lens.size(); i++) {
        if (lens[i].first == 1) {
            num++;
            if (num == K - 1) {
                swap(lens[i+1], lens[i+2]);
            }
        }
    }
    for (auto [val, num] : lens) {
        for (int i = 0; i < num; i++) cout << val;
    }
    cout << endl;
}