【AtCoder】ABC 393 D - Swap to Gather | 緑コーダーが解くAtCoder

atcoder.jp

配点: 425 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 533

問題概要
制約
考察
実装例
コメント

問題概要

0, 1からなる長さ $N$ の文字列 $S$ が与えられる。これから、以下の操作を0回以上行うことですべての1がひとかたまりになった状態にしたい。

操作 : $1\leq i\leq N-1$ を満たす整数 $i$ を選び、 $S_i$ と $S_{i+1}$ を入れ替える。

このとき、必要な操作回数の最小値を求めよ。

制約

$N$ は整数で、 $2 \leq N\leq 5\times 10^5$ を満たす。
$S$ には1が1つ以上含まれる。

考察

サンプルケース1や3などについて手を動かしながら考えていると、1の位置の中央値に、他の1を集めていけばよいのでは？というアイデアが思いつく。思いついちゃったんだから仕方がない...

ただ、未証明のままやるわけにもいかないので一応証明をしていく。以下では 0-indexed を前提とする。

$S$ 内の左から $i$ 個目の1のインデックスを $\mathrm{pos}_i$ 、1の個数( $\mathrm{pos}$ の長さ)を $c$ とする。また、最終的な $S$ が

$S_0$ から $S_{x-1}$ が0
$S_x$ から $S_{x+c-1}$ が1
$S_{x+c}$ から $S_{N-1}$ が0

となるように操作を行うとする。

ここで、1同士を入れ替えるという操作は意味がないことを踏まえると、開始時点で $S$ の左から $\mathrm{pos}_i$ 番目にあった1は、最終的に $S$ の左から $(\mathrm{x+i})$ 番目に移動していればよいことになる。

このときの操作回数 $f(x)$ は、

$\displaystyle f(x) = \sum_{i=0}^{c-1} |\mathrm{pos}_i - (x+i)| = \sum_{i=0}^{c-1} |(\mathrm{pos}_i - i) -x|$

と書くことができる。なお、数列 $\{ a_i \} \: (a_i= \mathrm{pos}_i - i)$ は単調増加数列であることに注意。

一般に、絶対偏差の和が最小となるのは $x = a_{\lfloor \frac{c}{2} \rfloor}$ (データの中央値)のときであるから(下記リンク参照)、求める最小操作回数は $f(a_{\lfloor \frac{c}{2} \rfloor})$ となる。

hiraocafe.com

(証明終)

ということで、あとはこの内容をもとに実装していこう。

実装例

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

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

using ll = long long;

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

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

    vector<int> pos;
    rep(i, 0, N)
    {
        if (S[i] == '1')
            pos.push_back(i);
    }

    int l = pos.size();
    rep(i, 0, l) pos[i] -= i;
    int mid_pos = pos[l / 2];

    ll ans = 0;
    rep(i, 0, l) ans += abs(pos[i] - mid_pos);

    cout << ans << endl;

    return 0;
}