🎉

ABC190 F - Shift and Inversions

2021/04/10に公開

Python

競技プログラミング

AtCoder

tech

解説

フェニック木
以下の手順で解く．

与えられた数列の転倒数をフェニック木で求める $O(NlogN)$
各 $k$ の時の転倒数を $k-1$ の時の転倒数から $O(1)$ で求める

まず，転倒数
ある数 $A_j(0\leqq j\leqq N)$ よりも前にある数 $A_i(0\leqq i\leqq j)$ の内で $A_j<A_i$ となる $A_i$ の個数が転倒数．
これは数列 $A$ を前から順に記録しながら見ていき，各 $A_j$ において $A_j$ よりも大きい数が記録内にいくつあるかを調べればよさそう．
これをフェニック木で求める．フェニック木についての詳しいことは参考資料とコード中のコメントで．
フェニック木は要素の更新と区間和の計算を $O(logN)$ で行える．
なので，木の要素を $0$ で初期化した後に各 $j(0\leqq j\leqq N)$ で以下の処理をすれば良さそう．

$A_j$ よりも前にある $A_j$ よりも大きい数の個数を転倒数に加算
$(A_jよりも前にあるA_jよりも大きい数の個数)=j-(A_jよりも前にあるA_jよりも小さい数の個数)$
木のインデックス $A_j$ をインクリメント
インデックス $A_i(0\leqq i\leqq j)$ が $1$ なら $A_i$ は既出， $0$ ならば既出でないということを表す

フェニック木なら上記の2つの操作はそれぞれ $O(logN)$ で行えるので，全体で $O(NlogN)$ で済む

続いて $k$ を変化させる．
各 $k$ で転倒数を求めるのにフェニック木を使うと，全体で $O(N^2logN)$ かかるので，別の方法が要る．
適当なサンプルで手計算してみると， $k$ 番目の解となる転倒数 $inv(k)$ は，

inv(k)=inv(k-1) + N - 2A_k + 1

で表せると分かる．
よって各 $k$ 当たり $O(1)$ で求まる．

コード

class fenwick:
    """Fenwick Tree (Binary Indexed Tree)

    フェニック木は部分和の更新と区間和の計算を効率的(O(logN))に行える木構造です．
    木の値を保持するリスト(table)の内，使用可能なインデックスは1以上です．
    これはLSBを用いて各頂点の親子関係を簡単に導出するためです．

    Attributes:
        table: 木の各頂点の値を記録

    Raises:
        IndexError: tableの使用していないインデックスへアクセスしようしています．\
            tableの使用できるインデックスは1以上です．
    """
    def __init__(self, init_iter: list = []):
        """木構造を生成します．

        リストの引数があれば，そのリストの要素で木が生成されます．
        引数が無い場合は，後々pushメソッドを用いて要素を追加していく必要があります．

        Args:
            init_iter (list, optional): 木の生成に用いられるリストです．
                デフォルト値は空リストです．
        """
        self.table = [0]
        for x in init_iter:
            self.push(x)

    def push(self, x: int = 0):
        """tableの末尾に要素xを追加し，必要な頂点を更新します．

        Args:
            x (int): 追加される要素xです．デフォルト値は0です．
        """
        num_elems = len(self.table)
        index = 1
        lsb = num_elems & -num_elems  # ここでのlsbは2の補数を表します
        while index != lsb:
            x += self.table[num_elems - index]
            index *= 2

        self.table.append(x)

    def sum(self, i: int) -> int:
        """[0,i]の範囲の区間和を求めます．

        Args:
            i (int): 区間和の末尾のインデックスです．半開区間では無く閉区間です．

        Raises:
            IndexError: クラスのdocstringを参照してください．
        """
        if i == 0:
            raise IndexError("インデックス0は使用できません．\
                詳しくはクラスのdocstringを参照してください")
        prefix_sum = 0
        while i > 0:
            prefix_sum += self.table[i]
            i -= i & -i
        return prefix_sum

    def add(self, i: int, x: int):
        """インデックスiの要素にxを加算します．

        Args:
            i (int): 加算される要素のインデックスです．
            x (int): 加算する値です．

        Raises:
            IndexError: クラスのdocstringを参照してください．
        """
        if i == 0:
            raise IndexError("インデックス0は使用できません．\
                詳しくはクラスのdocstringを参照してください")
        while i < len(self.table):
            self.table[i] += x
            i += i & -i


N = int(input())
A = list(map(int, input().split()))

bit = fenwick([0 for _ in range(N)])
ans = [0]
for i in range(N):
    bit.add(A[i] + 1, 1)
    ans[0] += i + 1 - bit.sum(A[i] + 1)

for a in A[:-1]:
    ans.append(ans[-1] + N - 2 * a - 1)

print(*ans, sep="\n")

解説

コード

参考

フェニック木

転倒数

Discussion