ABC190F - Shift and Inversions

問題概要

問題へのリンク

$0$から$N-1$までの整数をちょうど一つづつ含む数列$A$が与えられる。 $k \in \mathbb{Z}$に対して、数列$B$を次のように定める。 $$ \begin{equation*} B \coloneqq \{b_i\}_{i=0}^{N-1}, ~ b_i = a_{i+k ~ \mathrm{mod} ~ N} \end{equation*} $$ $k = 0, 1, 2, \dots , N-1$のそれぞれに対して、$B$の転倒数を求めよ。

制約

• $2 \leq N \leq 2 \times 10^5$

考察

$k$を一つ変える操作は、変える前の数列を左シフトしたようなものになっている。 例えば、$A = (0, 1, 2, 3, 4)$であったとき、

• $k=0$ : $B = (0, 1, 2, 3, 4)$
• $k=1$ : $B = (1, 2, 3, 4, 0)$
• $k=2$ : $B = (2, 3, 4, 0, 1)$
• $k=3$ : $B = (3, 4, 0, 1, 2)$
• $k=4$ : $B = (4, 0, 1, 2, 3)$

となる。 まずは転倒数について確認しておく。 転倒数とは、簡単に定義するなら、次の条件を満たす組$(i, j)$の個数である。

1. $i < j$
2. $B_i > B_j$

言葉で言うなら、各要素に対しての「自分より左にいる、自分より大きなものの総数」の和である。

さて、操作によって転倒数がどう変化するかを考えよう。 操作は、「一番左の要素を一番右に持っていく」というものになっている。 これを行うと、全体の転倒数は次のように変化することがわかる。

1. 動かした要素より小さい要素の数だけ転倒数が減る。
2. 動かした要素より大きな要素の数だけ転倒数が増加する。

これは転倒数の定義からも簡単に確認できるだろう。 これで、操作1回によって転倒数がどれだけ変化するかはすぐ算出できることがわかった。

後は最初の数列の転倒数を求められればACできそうだ。

典型テク: 転倒数を$O(N \log N)$で求める。

さて、数列$A$に対してナイーブに転倒数を求めるアルゴリズムを考えると、次のようなものになるだろう。

int InversionNumber = 0;
for (int i = 0; i < N; i++) {
    for (int j = 0; j < i; j++) {
        if (A[j] < A[i]) InversionNumber++;
    }
}

これは明らかに$O(N^2)$で、間に合わない。 さてどうしよう。

実は、転倒数を$O(N \log N)$で求めるアルゴリズムが存在する。 ここからその説明をする。

転倒数の定義から、次が成立する。 $$ \begin{equation*} (転倒数) = \sum_{i=0}^{N-1} ((j<i) ~ \land ~ (A_i < A_j)を満たすjの数) \end{equation*} $$ 少し寄り道をする。 まず、非負数列$B$に対して数列$C$を次のように定める。 $$ \begin{equation*} C \coloneqq \{c_i\}_{i=0}, ~ c_i = (数列Bに含まれるiの個数) \end{equation*} $$ 具体例を挙げよう。 数列$(0, 1, 2, 3)$に対して、$C = (1, 1, 1, 1, 0, 0, \dots)$であり、 数列$(1, 1, 3, 5, 7)$に対して、$C = (0, 2, 0, 1, 0, 1, 0, 1, 0, \dots)$となる。 また、数列$B$に対して定められる数列$C$を$C_B$と表すことにする。

数列$A$の$i$項目から$j$項目までを順番を変えずに切り取った数列を、$A$の連続部分数列と呼び、 $A[i \dots j]$と表すとする。

これらを用いて先程の転倒数の定義を言い換えると、次のようになる。 $$ \begin{equation*} (転倒数) = \sum_{i=1}^{N-1} \sum c \in ( C_{A[0 \dots i-1]}[A_i+1 \dots \infty] ) \end{equation*} $$ 式が最悪すぎるのには目をつぶってほしい…数学に明るくなく、うまい定式化ができなかった。

さて、ここで$C_{A[0 \dots i-1]}$については各ケースに対して1から求める必要は無く、 $C_{A[0 \dots i-2]}$に$A_i$を含めてやればよいことがわかるだろう。 また、数列$C$は(数列$A$の最大値)番目以降は全て0になるため、そこで打ち切って良い。 したがって、各ケースで 「数列の要素を一つだけ変更する。新しく得た数列の$A_i+1$番目から最後までの総和を求めよ」という問題を解けば良い。

この問題は、数列の動的区間和を高速に求めるデータ構造(BITやSegment Tree)を用いることで解ける。 例えば、Segment Treeを使った場合、区間和の更新と取得で$O(\log N)$かかるため、全体で$O(N \log N)$になる。

以上で最初の一つの転倒数を$O(N \log N)$で得ることができた。

補足: 本章では必要以上にややこしい説明を行っていますが、これは自分が後で見返したときのためです。 正直ググって他の資料をあたったほうがわかりやすいと思います。ごめんなさい。

実装例

import std;

void main () {
    int N = readln.chomp.to!int;
    int[] a = readln.split.to!(int[]);

    solve(N, a);
}

void solve (int N, int[] a) {
    /* 転倒数 O(NlogN) */
    auto RSQ = new SegmentTree!(int, (int a, int b) => (a+b), () => 0)(a.length);
    long[] ans = new long[](N);
    foreach (i; 0..a.length) {
        ans[0] += RSQ.prod(a[i]+1, a.length);
        RSQ.set(a[i], RSQ.get(a[i])+1);
    }

    /* 左端を右端に動かすとき、(もとの数列がわかれば)転倒数の変化は比較的簡単にわかるので、最初の一回を求められれば後はO(N) */

    foreach (i; 1..N) {
        ans[i] = ans[i-1] - a[i-1] + (N-a[i-1]-1);
    }

    foreach (i; 0..N) writeln(ans[i]);
}

Segment Treeの実装は長いので省略した。

感想

自力でACできたのと、転倒数を$O(N \log N)$で求めるアルゴリズムに触れたので良かった。 このように、「すでにわかっている結果から少しいじると他の解がわかる」 みたいなタイプの問題を落とさないようにしたい(3戦1勝2敗)。