実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 223

問題概要

0, 1のみからなる長さ $N$ の文字列 $S$ が与えられる。 $S$ の中で先頭から $K$ 番目の1の塊を、 $K-1$ 番目の1の塊の直後まで移動した文字列を出力せよ。なお、 $S$ には1の塊が $K$ 個以上含まれることが保証される。

制約

$1 \leq N \leq 5\times 10^5$
$2 \leq K$

考察

まずは0の塊と1の塊をデータとして持っておきたい。こういうときはランレングス圧縮が有効だ。pair<char, int>の配列のデータ構造に、それぞれの塊の大きさ (長さ？) を格納できる。

その後、0の塊と1の塊は互いに隣接している (どちらか一方が連続することはない) ことを利用して、先頭から $K$ 個目の1の塊とその直前の0の塊の位置をswapする。

あとは先頭から順に出力していけばよい。

コード

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

std::vector<std::pair<char, int>> encode(const std::string &str)
{
    int n = (int)str.size();
    std::vector<std::pair<char, int>> res;
    for (int l = 0; l < n;)
    {
        int r = l + 1;
        for (; r < n && str[l] == str[r]; r++)
        {
        };
        res.push_back({str[l], r - l});
        l = r;
    }
    return res;
}

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

    vector<pair<char, int>> rle = encode(S);
    int cnt = 0, idx = 0;
    while (cnt != K - 1)
    {
        if (rle[idx].first == '1')
            cnt++;

        idx++;
    }

    swap(rle[idx], rle[idx + 1]);

    for (auto &&p : rle)
    {
        rep(i, 0, p.second)
        {
            cout << p.first;
        }
    }
    cout << endl;
}