この記事は帰ってきた AOJ-ICPC Advent Calendar 2022の21日目の記事です。
この問題は AtCoder にも存在します。
問題概要
整数 $N,C$ と、長さ $N$ の順列 $b$ 、整数列 $w$ が与えられます。 順列 $p=(1,2,...,N)$ に以下の操作を好きな回数行い $p=b$ とするのにかかる最小のコストを出力してください。
操作:
$1\leq i\leq N-1$ を満たす $i$ を選び、 $p_{i},p_{i+1}$ を入れ替える。かかるコストは $\min(Cw_{p_{i}}+w_{p_{i+1}},w_{p_{i}}+Cw_{p_{i+1}})$
$1\leq N\leq 10^{5}$
$1\leq C\leq 100$
$1\leq w_{i}\leq 10^{5}$
解法
かかるコストが $1$ のときは $b$ の転倒数と同じです。
転倒数の求め方は最初に $ans=0,A_{1}=A_{2}=...=A_{N}=0$ として、$i=1,2,...,N$ の順に以下の操作を行うと求まります。
$ans\leftarrow ans+\sum_{k=1}^{b_{i}-1}A_{k}$ としたあと、 $A_{b_{i}}\leftarrow 1$ とする。
点更新区間和なので、セグメントツリーで $O(N\log N)$ 求まります。
$C=1$ のときはコストが $w_{p_{i}}+w_{p_{i+1}}$ になり、転倒数のときと同じようにセグメントツリーを用いて $O(N\log N)$ で求まります。
具体的には、$ans=0,A_{1}=A_{2}=...=A_{N}=0,B_{1}=B_{2}=...=B_{N}=0$ として、$i=1,2,...,N$ の順に以下の操作を行うと求まります。
$ans\leftarrow ans+\sum_{k=1}^{b_{i}-1}A_{k}+Cw_{i}\sum_{k=1}^{b_{i}-1}B_{k}$ としたあと、 $A_{b_{i}}\leftarrow w_{i},B_{b_{i}}\leftarrow 1$ とする。
$2\leq C$のときは、$w$ の大小によってコストの計算方法が異なります。以下の方法をセグメントツリーを用いれば $O(N^{2}\log N)$で求まります。
$ans=0$ 、 $N\times N$ の二次元配列 $A,B$ を用意します($A,B$ の初期値は全て $0$)。まず、「$w_{i}$ が $w$ の中で $X_{i}$ 番目に小さい」となるように順列 $X$ となるように定義します。($w_{i}$ が同じのときは適当に順序をつけて $X$ に同じ要素が存在するようにする。)
そして $i=1,2,...,N$ の順に以下のように操作を行う。
$ans\leftarrow ans+\sum_{j=1}^{b_{i}-1}(\sum_{k=1}^{X_{i}-1} (CA_{j,k}+B_{j,k}w_{i})+\sum_{k=X_{i}+1}^{N} (A_{j,k}+CB_{j,k}w_{i}))$
と更新した後
$A_{b_{i},X_{i}}\leftarrow w_{i},$ $B_{b_{i},X_{i}}\leftarrow 1$ と更新する。
これを高速化するのですが、二次元セグメントツリーと同等の考え方をすると $O(N (\log N)^{2})$ で解くことができます。
しかし、実装が多分大変だと思うので、平方分割を使います。
実装は綺麗にできていると思うので、詳しくは実装を見てほしいですが、簡単に説明します。
$w$ の大きさに応じて $D$ 個のグループに分けます。
自身 $i$ より $w$ が小さいグループに所属している $k$ に対して、 $k<i$ かつ $b_{k}<b_{i}$ なら $ans\leftarrow ans+Cw_{k}+w_{i}$
自身 $i$ より $w$ が大きいグループに所属している $k$ に対して、 $k<i$ かつ $b_{k}<b_{i}$ なら $ans\leftarrow ans+w_{k}+Cw_{i}$
と更新し、同じグループに対しては愚直に行います。
この実装の計算量は、一つの $i$ が一つのグループに対して、 $O(\log N)$、愚直パートが $O(N^{2}/D)$ なので、
合計で $O(\frac{N^{2}}{D}+ND\log N)$ です。
一つの $i$ が一つのグループに対して、 $O(\log N)$ で計算できるのは、index を二分探索で求め、 $w$ の和や $k<i$ を満たすものの数を前述した $C=1$ のパートのときのように求めることを考えると、 $i$ 毎の更新が $O(\frac{N}{D})$ でできます。(セグメントツリーを使わずに、配列で累積和を持っておく)
言語化するのが難しいです(何のためにブログを書いているんだ)。ぜひコードを見てほしいです。
コード
先ほど定義した変数名と違うかもしれませんがお許しください。
AOJ で 0.71 sec
AtCoder で 953ms
#include<bits/stdc++.h> using namespace std; #define rep(i,a,b) for(int i=a;i<b;i++) using ll = long long; #define all(p) p.begin(),p.end() int main(){ ll N,C; cin>>N>>C; vector<ll> B(N),W(N); rep(i,0,N) cin>>B[i]>>W[i],B[i]--; vector<ll> w_order(N); iota(w_order.begin(),w_order.end(),0); sort(all(w_order),[&](int l,int r){ return W[l]<W[r]; }); ll D=min(N,100ll); vector<int> pos(N); vector<int> L(D+1); rep(i,0,D+1){ L[i]=(N*i+D-1)/D; } rep(i,0,N){ int tmp=(i*D)/N; pos[w_order[i]]=tmp; } vector<vector<int>> A(D); vector<vector<int>> count(D); vector<vector<int>> sum(D); rep(i,0,N){ A[pos[i]].push_back(B[i]); } rep(i,0,D){ sort(all(A[i])); count[i].resize(A[i].size()+1,0); sum[i].resize(A[i].size()+1,0); } ll ans=0; rep(i,0,D){ rep(k,L[i],L[i+1]) rep(l,k+1,L[i+1]){ int a=w_order[k],b=w_order[l]; if((a<b)^(B[a]<B[b])){ ans+=C*W[a]; ans+=W[b]; } } } rep(i,0,N){ rep(j,0,D){ if(j==pos[i]) continue; int ind=lower_bound(A[j].begin(),A[j].end(),B[i])-A[j].begin(); if(j<pos[i]){ ans+=sum[j][ind]*C; ans+=count[j][ind]*W[i]; }else{ ans+=sum[j][ind]; ans+=count[j][ind]*W[i]*C; } } rep(k,0,(int)(A[pos[i]].size())){ sum[pos[i]][k]+=W[i]; count[pos[i]][k]++; if(A[pos[i]][k]==B[i]) break; } } cout<<ans<<"\n"; }
終わりに
二次元セグメントツリー系は全部こういう感じにできるんですかね?
帰ってきた AOJ-ICPC Advent Calendar 2022の僕の担当パートはこれで最後です。
来週の水曜日は横浜2022です。頑張りたい。