Snackdown 2017 Pre-elimination A: Consecutive Snakes

これは良問だなあと思った。

問題概要

長さ $L (1 \leq L \leq 10^{9})$ の蛇が数直線上に $N (1 \leq N \leq 10^{5})$ 匹おり、 $i$ 番目の蛇の左端の座標は $S_{i} (1 \leq S_{i} \leq 10^{9})$ である。つまり、 $i$ 番目の蛇は区間 $\left[ S_{i}, S_{i} + L \right]$ にいる。

いま、この蛇たちを以下の条件が成り立つように移動させたい。

N 匹の蛇すべてについて、区間 $\left[ A, B \right] (1 \leq A, B \leq 10^{9})$ の中にいる
左から $i$ 番目の蛇の右端と $i+1$ 番目の蛇の左端は繋がっている。つまり $1$ 番目の蛇が区間 $\left[ X, X + L \right]$ にいるとき、 $2$ 番目の蛇は $\left[ X + L, X + 2L \right]$ におり、 $\dots$ 、 $N$ 番目の蛇は $\left[ X + (N-1)L, X + NL \right]$ にいる。

ある蛇が $\left[ X_{1}, X_{1} + L \right]$ から $\left[ X_{2}, X_{2} + L \right]$ に移動するとき、コストが $| X_{1} - X_{2} |$ かかる。全体コストはこのコストを $N$ 匹に対して計算した時の総和である。条件が成り立つように移動させた時の全体コストの最小値を出力せよ。

解説

まず、条件より以下のことがわかります。

一番左にくる蛇の左端の座標を決定すると、全ての蛇の左端の座標が決まる
移動前に左から $i$ 番目にいた蛇は、移動後も左から $i$ 番目にいたほうが良い

つまり、蛇を座標順にソートしておいて、一番左の蛇がどこにいるか決めてしまえば、コスト計算はできるということです。しかし、制約を見ると一番左の蛇の配置パターンは非常に多いことがわかるので、何か工夫が必要です。そこで、一旦定式化して考えましょう。

一番左の蛇の左端を $x$ とします。すると、 $i$ 番目の蛇の左端は $x + i L$ と表せます (0-indexed)。移動前に左から $i$ 番目にいた蛇の左端を $S_{i}$ とすると、全体コストは以下のように書けます。

\begin{align} cost &=& \sum_{i=0}^{N-1} | S_{i} - (x + iL) | \\ &=& \sum_{i=0}^{N-1} | (S_{i} - iL) - x | \\ &=& \sum_{i=0}^{N-1} | T_{i} - x | \end{align}

$S_{i} - iL$ は $x$ の値の設定に関係なく決まるので、定数とみなして良いです。これを新たに $T_{i}$ とします。

$T_{i}$ をソートすると、 $T_{0} \leq T_{1}\leq \dots \leq T_{N-1}$ のような形になります。さて、 $T_{i} \leq x$ を満たす最大の $i$ を $j$ とすると、全体コストはこのように書けます。

$cost = \sum_{i=0}^{j} (x - T_{i}) + \sum_{i=j+1}^{N-1} (T_{i} - x)$

コスト関数は絶対値付きの一次関数を足しあわせた形になっているので凸関数です。よって、 $x$ について二分探索して解くことができます。

計算量はソート $O(N \log N)$ 、二分探索パートでコスト計算時に upper_bound が必要なため $O(\log N \log C)$ です $(C = B - A)$ 。

ソースコード

累積和を 0-indexed にすると配列外参照が起こるので、そこだけ 1-indexed で書いています。

int T, N, L, A, B;
int s[100010];
int t[100010], acc[100010];

int get(int l, int r) {
    l++; r++;
    return acc[r] - acc[l-1];
}

int calc(int x) {
    int idx = upper_bound(t, t+N, x) - t;
    int vl = idx * x - get(0, idx-1);
    int vr = get(idx, N-1) - (N-idx) * x;
    return vl + vr;
}

signed main() {
    scanf("%lld", &T);
    rep(cases, 0, T) {
        scanf("%lld%lld%lld%lld", &N, &L, &A, &B);
        rep(i,0,N) scanf("%lld", &s[i]);
        sort(s, s+N);

        // 累積和は 1-indexed
        rep(i,0,N) t[i] = s[i] - i * L;
        sort(t, t+N);
        rep(i,0,N) acc[i+1] = acc[i] + t[i];

        // A <= x <= B を満たす x で最小を求めよう
        int lb = A - 1, ub = B - N * L;
        while(ub - lb > 1) {
            int mid = (ub + lb) / 2;
            if(calc(mid) < calc(mid+1)) ub = mid;
            else lb = mid;
        }
        cout << calc(ub) << endl;
    }
    return 0;
}