AOJ 2853 : Librarian's Work (jag 夏合宿2017:day3I)

この記事は帰ってきた AOJ-ICPC Advent Calendar 2022の21日目の記事です。

問題のリンク

この問題は AtCoder にも存在します。

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です。頑張りたい。