【競プロ】競プロ典型90問 066 - Various Arrays（★5）で想定解よりかなり早く解く

はじめに

こんにちは！！競技プログラミングをやっている高校生のButterFlvです！この記事では競プロ典型90問 066 - Various Arrays（★5）の問題で想定解として解説のある $O((NM)^2)$ （$M=\mathrm{max}(L_i,R_i)$ とおいた.）の解法を改善した $O(M+N\log M)$ の時間計算量で解くことができたので紹介していきます。

誤り等あればコメントなどで指摘していただけるとありがたいです！！
またこの記事は問題の内容に関するネタバレを含みます。

問題

問題文

数列屋の高橋君は長さ $N$ の数列 $a$ を作っています。数列 $a$ の $i$ 番目の要素 $a_i$ は、$L_i$ 以上 $R_i$ 以下の整数から一様ランダムに選ぶことで決定されます。

このようにしてできた数列 $a$ の転倒数の期待値を求めてください。

なお、長さ $m$ の数列 $x$ の「転倒数」とは、$i\lt j$ かつ $x_i\gt x_j$ であるような $(i,j)\ (1\le i,j\le m)$ の個数のことです。

制約

• $1\le N\le 100$
• $1\le L_i\le R_i\le 100$
• 入力はすべて整数で与えられる

ステップ０（愚直解）

数列としてありうるものの総数が $\prod(R_i-L_i+1)$ であることからこれら一つ一つに対して $O(N^2)$ で転倒数を求めればよく、この解法は $O(N^2M^N)$ です。（$M=\mathrm{max}(L_i,R_i)$ とおいた.）本問の制約においてこれは十分低速で、最悪 $O(N^2M^N)\approx 100^2\cdot 100^{100}=10^{204}$ となります。やはりアルゴリズムは素晴らしいですね～

ステップ１（想定解）

期待値の線形性を利用して $(i,j)\ (1\le i\lt j\le N)$ の組それぞれに対して独立に「転倒している」（転倒数に$+1$の寄与をする）期待値を求め、最後に足しあわせるという解法です。$a_i, a_j$ の値を $[L_i,R_i],[L_j,R_j]$ で独立に動かしているので計算量は $O((MN)^2)$ です。最悪 $O((MN)^2)\approx (100\cdot 100)^2=10^8$ で、これに $\dfrac{1}{2}$ や $\dfrac{1}{4}$ の定数倍がつくので十分高速です。

ステップ２（dp(平面走査)でやってみる）

まずは、すべての $i$ で $L_i=R_i$ のときを考えてみましょう。これは素直に $a_i=L_i(=R_i)$ と、数列そのものが与えられている状況と同じだといえるので、BITやセグ木を使った転倒数を求めるコード（$\cdots$★）を書いてあげればよいです。これと同じようなことを $L_i\lt R_i$ のときにもできないか考えてみます。（★）と同様に列を左から一つずつ見ていきます。新たに列の右に $i$ 項目を追加するとき、どのような情報が配列dpにあればうれしいかを考えます。それは、整数 $x$ について、$a$ の $i-1$ 項目までに $x$ が含まれる個数の期待値です。（この考え方は（★）に適用しても列が確定しているので $x$ が含まれる個数の期待値がきっちり整数になっていると解釈でき、一般化として自然です。）よって、$j\in [L_i,R_i]$ とループする中で配列dpの $j$ より添え字が大きいものの総和をとり、さいごに $j\in [L_i,R_i]$ について足し合わせ、$R_i-L_i+1$ で割ったものを答えの変数に足してあげればよいです。また、ここで配列dpの $L_i$ 項目から $R_i$ 項目までにそれぞれ $\dfrac{1}{R_i-L_i+1}$ を加算すれば列に $i$ 項目を追加できたことになります。上記をそのまま実装すると $O(NM^2)$ です。$j$ を固定する平面走査で少し高速化できました。

ステップ３（dpの高速化）

ここから適切なデータ構造を用いてさらに高速化していきます。ステップ２では列を左からひとつづつ見ていくときに区間の総和を何回もとり、また最後に区間に加算していました。これは遅延評価セグメントツリーでできそうです。実際、$j\in [L_i,R_i]$ について区間 $[j+1,M]$ で総和を取得することでこの操作を $O(M\log M)$ に高速化でき、最後の区間加算は遅延評価セグメント木の得意な操作ですので $O(\log M)$ にできます。これで $O(NM\log M)$ です。しかし、ここからのひと工夫で $O(M+N\log M)$ まで落とせます。暫定の解法でもっともネックな部分は $j\in [L_i,R_i]$ について区間 $[j+1,M]$ で総和を取得している部分です。これが足されている様子を観察すると、最終的には階段状に足されていることがわかります。厳密にいえば、dp配列の各項について操作によってそれぞれが、何回総和に寄与しているかを考えると、$\mathrm{dp}[L_i+j]\ (0\le j\le R_i-L_i)$ は $j+1$ 回足され、$j\gt R_i-L_i$ を満たす $j$ について $\mathrm{dp}[L_i+j]$ は $R_i-L_i+1$ 回足されているとわかります。これは遅延セグメント木に $\sum_{1\le j\le N}\mathrm{dp}[j]$ と $\sum_{1\le j\le N}j\ \mathrm{dp}[j]$ の情報を持たせれば適宜足し引きすることで $j\in [L_i,R_i]$ のループを回避して一発で求められるため、取得する操作も $O(\log M)$ にできます。セグ木の構築 $O(M)$ と合わせると全体の計算量は $O(M+N\log M)$ となりかなり高速化できました。

実装

Pythonでの実装

C++での実装

main.cpp

#include<bits/stdc++.h>
#include<atcoder/lazysegtree>
using namespace std;
using namespace atcoder;

struct S{ int len,wei; double sum,val; };
using F=double;

S op(S l, S r){ return S{l.len+r.len, l.wei+r.wei, l.sum+r.sum, l.val+r.val}; }
S e(){ return S{0, 0, 0.0, 0.0}; }
S mapping(F l, S r){ return S{r.len, r.wei, r.sum+l*r.len, r.val+l*r.wei}; }
F composition(F l, F r){ return l+r; }
F id(){ return 0.0; }

const int M=100;
int N;

int main(){
  cin>>N;

  vector<S> init;
  for(int i=0;i<=M;i++)init.push_back(S{1, i, 0.0, 0.0});
  lazy_segtree<S,op,e,F,mapping,composition,id>seg(init);

  double ans=0;
  for(int i=0;i<N;i++){
    int L,R; cin>>L>>R;
    double interval=R-L+1;
    S L_R_res=seg.prod(L+1, R+1);
    ans += (L_R_res.val - L_R_res.sum*L)/interval;
    ans += seg.prod(R+1, M+1).sum;
    seg.apply(L, R+1, 1/interval);
  }

  cout<<fixed<<setprecision(15);
  cout<<ans<<endl;
}

実装上の工夫

• 末端の各ノードに長さと重さという変数を持たせて遅延セグ木に乗るようにした。

さいごに

• 考察がとても楽しかったです！！
• ここまで計算量を落とせると思いませんでした。
• おそらく定数倍高速化を頑張れば制約が $1\le N\le 2\times 10^5$ と $1\le L_i\le R_i\le 2\times 10^5$ でも間に合うと思います。（要検証）