Atcoder Beginner Contest (ABC) 221 E - LEQ

問題文

想定解法はBIT(Binary Indexed Tree)でした。最初からBITだと分かって解いたというよりは、手元でノートに色々書き込みながら考察を進めていく中でなんとなく「BIT使えそうだな」→「あ、本当に行けそう」と徐々に確信が強くなり、解くことができました。以下では、本番で私が考察した順番でなるべく丁寧にステップを踏んで説明しています。長々と書いているので、結論だけを見たい方は正しい操作手順以降をご覧ください。

部分列の性質：左端と右端が決まればあとは自由

例があった方がわかりやすいので、以下の数列で計算することを考えてみましょう。

A = [4, 2, 7, 1, 5, 6, 3]

まず気づくのは、ここから部分列A'を選んだとき、問題文より部分列が正しい部分列として成り立つ条件はA_1' <= A_k'なので、部分列の左端と右端が決まればそれ以外は何でもよいということです。

例えば、A'の左端としてA_1 = 4を、右端としてA_5 = 5を選んだ場合を考えてみます。この時、左端と右端に挟まれたA_2 = 2, A_3 = 7, A_4 = 1はいずれも選んでも選ばなくても部分列の条件を満たします。つまり、

A' = [4, 5]（何も選ばない場合）
A' = [4, 2, 5]（A_2を選んだ場合）
A' = [4, 7, 5]（A_3を選んだ場合）
...
A' = [4, 2, 7, 5]（A_2, A_3を選んだ場合）
...
A' = [4, 2, 7, 1, 5]（全て選んだ場合）

のいずれもが正しい部分列として成り立ちます。左端と右端以外の要素１つ１つについて選ぶor選ばないの選択ができるので、左端の位置をl, 右端の位置をrとすると、あり得る部分列の個数は2r-l-1となります。

計算をまとめたい：右端を固定

次に、計算をまとめることを考えます。

例えば上の例

A = [4, 2, 7, 1, 5, 6, 3]

で、右端をA_5 = 5に固定して考えてみます。この時、左端の位置を1<=l<=4の範囲で移動させてみると、あり得る部分列の個数はそれぞれ

l = 1: 25-1-1 = 8通り
l = 2: 25-2-1 = 4通り
l = 3: 0通り（A_3 = 7 > 5で正しい部分列の条件を満たさないため除外）
l = 4: 25-4-1 = 1通り

となります。右端が左端より大きい場合を除外した上で、左端の位置によって場合の数が変わってくることがわかります。

これを一般化して、右端A_rを固定して考えたとき、うまくまとめて計算するには、①既出の数字(A_i, i < r)のうち、A_r以下のものの、②インデックスの位置i(i<r)を踏まえた場合の数の総和の２点が素早く分かればいいということになります。

転倒数の計算に使うBITとの共通点

①既出の数字(A_i, i < r)のうち、A_r以下のものをまとめて計算するのは、BITを使った転倒数の計算と似ています（ただし、転倒数はA_rより大きいものをまとめて計算する）。BITを使った転倒数の計算の詳細については例えば以下の記事を参照してください。

いかたこのたこつぼ　転倒数

以下では、BITのインスタンスをbitとして、bitの配列の左からi番目の要素にxを加える操作をbit.add(i, x), 左からl番目からr番目までの要素の総和を計算する操作をbit.query(l-1,r)（累積和の計算と同じで、左側が開になる半開区間）と書くことにします。

BITを使って、例えば以下のようなことができそうです。

Aを左の要素から順に見ていくことを考えます。BITの配列は最初、全て0で初期化されています。

BIT: [0, 0, 0, 0, 0, 0, 0]

まず最初にA_1=4の場所(1-indexed)に1を加算します（bit.add(4, 1)）。

（A_1 = 4を加算） BIT: [0, 0, 0, 1, 0, 0, 0]

続いて、A_2 = 2です。2の場所に1を加算する前に、BITの配列を1<=i<=2の範囲の合計値をチェックします（bit.query(0, 2)）。この場合は全て0（＝既出の2以下の要素の個数は0）なので、何も起こりません。続いて、２の場所に１を加算します（bit.add(2, 1)）。

（A_2 = 2を加算） BIT: [0, 1, 0, 1, 0, 0, 0]

続いて、A_3 = 7です。7の場所に1を加算する前に、BITの配列を1<=i<=7の範囲の合計値をチェックします（bit.query(0, 7)）。この場合は２です。つまり、既出の要素の中で７以下の要素の個数は２ということがわかります。

しかし、これだと単に既出の要素の個数をまとめて計算しているだけで、②インデックスの位置i(i<r)を踏まえた場合の数の総和の計算ができていません。

インデックスの位置の情報を入れたい：BIT配列に加算する際に位置の情報を入れる

上の例では、加算の際に毎回１を加算していました。加算する１の代わりに、要素のインデックスの位置iの情報を入れたらどうなるでしょうか。

再び、Aを左の要素から順に見ていくことを考えます。BITの配列は最初、全て0で初期化されています。

BIT: [0, 0, 0, 0, 0, 0, 0]

まず最初にA_1=4の場所(1-indexed)にA_1のインデックスである1を加算します（bit.add(4, 1)）。これは上の例と同じです。

BIT: [0, 0, 0, 1, 0, 0, 0]

続いて、A_2 = 2です。まず加算する前に、BITの配列を1<=i<=2の範囲の合計値をチェックします（bit.query(0, 2)）。この場合は全て0（＝既出の2以下の要素の個数は0）なので、何も起こりません。これも、上の例と同じです。続いて、２の場所にA_2のインデックスの２を加算します（bit.add(2, 2)）。

BIT: [0, 2, 0, 1, 0, 0, 0]

続いて、A_3 = 7です。まず加算する前に、BITの配列を1<=i<=7の範囲の合計値をチェックします（bit.query(0, 7)）。この場合は2+1=3が計算結果として出てきます。

ここで、A_1 = 4については１、A_2=2については２がBIT配列に代入されており、近い要素ほど値が大きくなってしまっていることに気づくと思います。一方、今やりたいことは最初に見た

l = 1: 25-1-1 = 8通り
l = 2: 25-2-1 = 4通り
l = 3: 0通り（A_3 = 7 > 5で正しい部分列の条件を満たさないため除外）
l = 4: 25-4-1 = 1通り

のように、近い要素ほど値を小さくしたいのです。

近い要素ほど値を小さくしたい：マイナスにしてみる

上の例で加算するインデックスの位置情報にマイナスをかけることで、やりたいことに一歩近づく気がします。やってみましょう。

BITの配列は最初、全て0で初期化されています。

BIT: [0, 0, 0, 0, 0, 0, 0]

まず最初にA_1=4の場所(1-indexed)にA_1のインデックスである1にマイナスをかけた-1加算します（bit.add(4, -1)）。

BIT: [0, 0, 0, -1, 0, 0, 0]

続いて、A_2 = 2です。まず加算する前に、BITの配列を1<=i<=2の範囲の合計値をチェックします（bit.query(0, 2)）。この場合は全て0（＝既出の2以下の要素の個数は0）なので、何も起こりません。これも、上の例と同じです。続いて、２の場所にA_2のインデックスにマイナスをかけた-2を加算します（bit.add(2, -2)）。

BIT: [0, -2, 0, -1, 0, 0, 0]

続いて、A_3 = 7です。まず加算する前に、BITの配列を1<=i<=7の範囲の合計値をチェックします（bit.query(0, 7)）。この場合は-2+(-1)=-3が計算結果として出てきます。

確かに、近い要素ほど値が小さくなってはいます。しかし、求めたいものは2r-l-1で、２冪の形になっています。

２冪の形にしてみる

上の例からさらに１歩進んで、２x（xはマイナスをかけたインデックスの値）を加算することを考えてみます。

BITの配列は最初、全て0で初期化されています。

BIT: [0, 0, 0, 0, 0, 0, 0]

まず最初にA_1=4の場所(1-indexed)に2-1を加算します（bit.add(4, pow(2, -1))）。

BIT: [0, 0, 0, 2-1, 0, 0, 0]

続いて、（合計値のチェックは0なので省略）A_2=2の場所に2-2を加算します（bit.add(4, pow(2, -2))）。

BIT: [0, 2-2, 0, 2-1, 0, 0, 0]

続いて、A_3 = 7です。まず加算する前に、BITの配列を1<=i<=7の範囲の合計値をチェックします（bit.query(0, 7)）。この場合は2-2 + 2-1が計算結果として出てきます。

この結果は正しいでしょうか？

本来であれば、右端をA_3 = 7で固定した場合、ありうる部分列の個数は、左端lを1<=l<=2の範囲で動かしてみると、

l=1: 23-1-1 = 2通り
l=2: 22-1-1 = 1通り

とならないといけません。

毎回、計算結果に２冪をかけて値を調整する

よく観察すると、以下のように、求めたい結果は、BITでの計算結果に同じ２冪の値をかけたものと等しいことがわかります。

  • l=1
    • BITの計算結果：2-2
    • 求めたい結果：21（BITの計算結果に23をかけたもの）
  • l=2
    • BITの計算結果：2-1
    • 求めたい結果：22（BITの計算結果に23をかけたもの）

つまり、BITに入力された値は２の「インデックス（絶対位置）にマイナスをかけた値」乗でしたが、求めたい値は左端と右端の位置関係（相対位置）によって決まるので、毎回右端の絶対位置に対応した２冪の値をかけて相対位置に対応した値に変換してあげる必要があるのです。

mod逆元の利用

注意点として、この問題では計算結果のmodをとった値が求められているので、上の計算結果のうち２の負冪（例えば2-2）についてはmod逆元を使った形で値を持っておく必要があります。

mod逆元の計算方法については例えば記事を参照してください。

正しい操作手順

以上を踏まえて、正しい操作手順は以下になります。まずAは以下です。

A = [4, 2, 7, 1, 5, 6, 3]

BITの配列は最初、全て0で初期化されています。答えとしてans = 0に初期化します。

BIT: [0, 0, 0, 0, 0, 0, 0]

まず最初にA_1=4の場所(1-indexed)に2-1 (mod p)を加算します（bit.add(4, pow(2, -1, p))）。

BIT: [0, 0, 0, 2-1, 0, 0, 0]

続いて、A_2 = 2です。まず加算する前に、BITの配列を1<=i<=2の範囲の合計値をチェックします（bit.query(0, 2)）。この場合は全て0（＝既出の2以下の要素の個数は0）なので、何も起こりません。続いて、BIT配列のA_2=2の場所に2-2 (mod p)を加算します（bit.add(2, pow(2, -2, p))）。

BIT: [0, 2-2, 0, 2-1, 0, 0, 0]

続いて、A_3 = 7です。まず加算する前に、BITの配列を1<=i<=7の範囲の合計値をチェックします（bit.query(0, 7)）。計算結果として2-2+2-1が返ってきますが、これに現在位置３に対応する23をかけて調整し、23(2-2+2-1) = 21 + 22を得ます。これをansに加算します。続いて、BIT配列のA_3=7の場所に2-2 (mod p)を加算します（bit.add(7, pow(2, -3, p))）。

BIT: [0, 2-2, 0, 2-1, 0, 0, 2-3]

続いて、A_4 = 1です。まず加算する前に、BITの配列を1<=i<=1の範囲の合計値をチェックします（bit.query(0, 1)）。1は一番小さい値なので計算結果として0が返ってきますが、24をかけても0なので、何も起こりません。続いて、BIT配列のA_4=1の場所に2-4 (mod p)を加算します（bit.add(1, pow(2, -4, p))）。

BIT: [2-4, 2-2, 0, 2-1, 0, 0, 2-3]

続いて、A_5 = 5です。まず加算する前に、BITの配列を1<=i<=5の範囲の合計値をチェックします（bit.query(0, 5)）。計算結果として2-4+2-2+2-1が返ってきますが、これに現在位置５に対応する25をかけて調整し、25(2-4+2-2+2-1) = 21 + 23 + 24を得ます。これをansに加算します。続いて、BIT配列のA_5=5の場所に2-5 (mod p)を加算します（bit.add(5, pow(2, -5, p))）。

BIT: [2-4, 2-2, 0, 2-1, 2-5, 0, 2-3]

この調子で最後まで計算します。

Aに同じ数字が複数含まれる場合も問題なし

なお、最終的な答えansは単にBITの計算結果（＝要素の総和）のさらに総和をとっているだけなので、例えばAに同じ数字が２つ以上含まれていて、BITの配列の同じ箇所に複数の値の合計値が格納されている場合、例えば

A = [4, 4, 7, 1, 5, 6, 3]

で、２回目の操作で

BIT: [0, 0, 0, 2-1+2-2, 0, 0, 0]

となった場合でも、正しい計算結果が得られます。

提出コード（PyPy）

BITはyaketakeさんの実装をお借りしています。

Pythonでは組み込み関数powを利用して2-i (mod p) （ただしpは素数）はpow(2, -i, p)として計算できますが、PyPyの場合はpowの第二引数に負数を持ってくることができないので、フェルマーの小定理 a^{p-1} \equiv 1 を利用してpow(2, p-i-1, p)として計算しています。

class BIT:
    '''https://tjkendev.github.io/procon-library/python/range_query/bit.html'''
    def __init__(self, n):
        self.n = n
        self.data = [0]*(n+1)
        self.el = [0]*(n+1)
    def sum(self, i):
        s = 0
        while i > 0:
            s += self.data[i]
            i -= i & -i
        return s
    def add(self, i, x):
        self.el[i] += x
        while i <= self.n:
            self.data[i] += x
            i += i & -i
    def query(self, i, j=None):
        if j is None:
            return self.el[i]
        return self.sum(j) - self.sum(i)
      
N=int(input())
*A,=map(int,input().split())
mod = 998244353

# 座標圧縮
v2i = {v:i+1 for i,v in enumerate(sorted(set(A)))}
B = [v2i[a] for a in A]

bit = BIT(N+1)
ans = 0
for i,b in enumerate(B):
  if i:
    res = bit.query(0,b)
    ans += res*pow(2,i-1,mod)
    ans %= mod
  bit.add(b,pow(2,mod-i-1,mod))
  
print(ans)