Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【yukicoder No.877 Range ReLU Query】2つの汎用的な解法

いかにも競プロ臭のする問題。感動した。

No.877 Range ReLU Query - yukicoder

問題

長さNの数列Aが与えられる。

長さQのクエリ(l,r,x)が与えられる。

各クエリについて、[l,r]の範囲でmax(ai-x, 0)の和を計算せよ。

制約: N<=105, Q<=105

方針1: beet法

ライターによる想定解法。サンプルを利用して説明する。

1333 133 (100) (33) 13 3 3 (3) (1)

と数列とクエリをミックスして降順ソートする。こうして、例えば、(100)の段階では

0 0 0 133 1333

であったと考えることが出来れば、[l,r]の区間和と、[l,r]間の非ゼロ数の個数がわかれば解くことが出来る。後者は、数値を入れたところを1とした区間和を別途持てばよい。

maxをとった時の0というのを「値がまだ入ってない」と考えることで、この解法に至ることが出来る。

この、大きい方からとっていくと、過去のものは全員今のやつより大きいという考え方は、

www.rustforbeginners.com

でも使われていて、汎用的だと思う。

大きい方からとっていく、後ろから見ていくは競プロではよくある考え方。

方針2: kmjp法

別解。

yukicoder : No.877 Range ReLU Query - kmjp's blog

エッセンスだけ話すために、平方分割で考える。

この問題は、仮にAが昇順ソートされていると二分探索と累積和で終了する。

Aを平方分割し、各バケットに、昇順ソートした配列と、その累積和を予め計算しておく。このコストは、root(N) root(N) log(root(N)) = N log(root(N))である。

こうすると、区間[l,r]について1バケットあたりlog(root(N))とバケット数root(N)で答えが計算出来る。

だからトータルでは、Q root(N) log(root(N))程度となり、めでたくTLEとなる。

TLEしないためには階層を深くする必要があり、それがkmjp法ということになる。

実装

方針1 beet法を実装した。

fn solve() {
    input!{
        N:usize,Q:usize,
        a:[i64;N],
        q:[(usize,usize,usize,i64);Q],
    }
    let mut ans=vec![0;Q];
    let mut v=vec![];
    for i in 0..N {
        v.push((a[i],1,i));
    }
    for i in 0..Q {
        v.push((q[i].3,0,N+i));
    }
    v.sort(); v.reverse();

    let mut sum=BIT::new(N);
    let mut cnt=BIT::new(N);
    for i in 0..v.len() {
        let (x,_,j) = v[i];
        // dbg!(x,j);
        if j>=N { // query
            let k=j-N;
            let (_,l,r,_)=q[k];
            let l=l-1; // 0オリ
            let r=r-1;
            let s = sum.sum(r+1)-sum.sum(l);
            let t = cnt.sum(r+1)-cnt.sum(l);
            ans[k]=s-t*x;
        } else { // add value
            // dbg!(j+1);
            sum.add(j+1, x);
            cnt.add(j+1, 1 as i64);
        }
    }
    for x in ans {
        println!("{}",x);
    }
}