AtCoder Grand Contest 006 C - Rabbit Exercise

http://agc006.contest.atcoder.jp/tasks/agc006_c

問題概要

1~Nの番号が振られたN匹のうさぎが一次元の数直線上に存在し、各うさぎの初期座標を示す数列Xが与えられる。うさぎi(2<=i<=N-1)がジャンプを行うと、うさぎ(i-1)かうさぎ(i+1)のどちらかが等確率で選ばれ、うさぎiは選ばれた方のうさぎに関して対称な座標に移動する。ここでうさぎの番号を要素にもつM要素の数列Aが与えられ、うさぎたちがこの順番にジャンプすることを1セットとする。ジャンプをKセット行うとき、各うさぎの最終的な座標の期待値を求めよ。

3 <= N <= 105

1 <= M <= 105

1 <= K <= 1018

解法

うさぎaが(a+1)の方にジャンプするとジャンプ後の座標は2X[a+1] - X[a]になり、(a-1)の方にジャンプすると2X[a+1] - X[a]になるから、1回のジャンプ後のaの座標の期待値は2つの平均をとってX[a+1] + X[a-1] - X[a] となる。この計算を繰り返していけば(実行時間はともかく)答えは求まる。あとは高速化が問題で、Kの大きさから行列累乗をしたくなるがO(M3logK)とかになってしまって無理。

ここでの正解は階差数列で考えることである。Xの階差数列Yを取ると、ジャンプの操作は以下のようになる。

  • X: X[i-1], X[i], X[i+1] → X[i-1], X[i-1] - X[i] + X[i+1], X[i+1]
  • Y: X[i] - X[i-1], X[i+1] - X[i] → X[i+1] - X[i], X[i] - X[i-1]

これをみると階差数列の方では前後の値が単にswapされていることがわかる。つまりジャンプ操作は階差数列の2要素のswapであるとみなせるようになる。このswapを1セット分シミュレートすれば、1セットでどの要素がどこに移動するかがわかる。これは要するに置換なので、あとはK回分の置換の積を繰り返し2乗法の要領で求めていけばO(MlogK)ですべてのジャンプが完了したあとの階差数列を得ることができる。ここで元の数列において先頭のうさぎの座標が変わることはないので、最後は階差数列の要素を順に足していくだけで元の数列を復元することができる。

感想

難しい問題あるある、1個だけあります

♪~~~~~♪~~~~~~~

階差とりがち

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
long[] perm_prod(const long[] A, const long[] P) {
    auto ret = A.dup;
    foreach (i; 0..A.length) {
        ret[i] = A[P[i].to!int];
    }
    return ret;
}
 
long[] perm_pow(const long[] A, long x) {
    long[] P = A.dup;
    long[] ret = new long[](P.length);
    foreach (i; 0..P.length) ret[i] = i;
    while (x) {
        if (x % 2) ret = perm_prod(ret, P);
        P = perm_prod(P, P);
        x /= 2;
    }
    return ret;
}
 
void main() {
    auto N = readln.chomp.to!int;
    auto X = readln.split.map!(to!long).array;
    auto s = readln.split.map!(to!long);
    auto M = s[0].to!int;
    auto K = s[1];
    auto A = readln.split.map!(x => x.to!int - 1).array;
 
    auto D = new long[](N - 1);
    foreach (i; 0..N-1) D[i] = X[i + 1] - X[i];
 
    auto B = new long[](N - 1);
    foreach (i; 0..N-1) B[i] = i;
    foreach (i; 0..M) swap(B[A[i] - 1], B[A[i]]);
 
    B = perm_pow(B, K);
    D = perm_prod(D, B);
 
    auto ans = new long[](N);
    ans[0] = X[0];
    foreach (i; 0..N-1) ans[i + 1] = ans[i] + D[i];
 
    ans.each!writeln;
}