Rustコトハジメ

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

【ABC140-E Second Sum】BTreeSetを使って実装する (1.17以上限定)

E - Second Sum

問題

長さNの順列Pが与えられる。

このうち、範囲[l,r]について2番目に大きいものをXlrとする。

すべての範囲について、Xlrを足し合わせた合計を求めよ。

方針

二番目に大きいというのがどういうことなのか考える。これは、あるPiを含む範囲を考えた時、「選択した範囲に自分より大きなものが一つだけある」ということである。

Piを大きな方から辿っていき、その都度インデックスを集合に入れていく。あるPjに来た時に、その集合のうち左にあるものと右にあるものがわかると、Pjが2番目に大きくなる範囲のとり方というのは、図における赤同士をとった時か、青同士をとった時ということになる。

f:id:akiradeveloper529:20190908181047j:plain

初見で思いつくのはかなり難しいと思う。おれも解説を理解するのはだいぶ時間がかかった。

実装

C++ならsetを使うと出来るようだが、Rustの場合はBTreeSetを使うといける。

range関数は、1.17からしか使えないので2019年9月8日時点のAtCoderでは使用することは出来ないが、これを使うと、その範囲の中にある値を返してくれるようである。

use std::collections::BTreeSet;
fn main()
{
    let mut s=BTreeSet::new();
    for i in 1..100 {
        s.insert(i);
    }
    let mut it = s.range(50..);
    dbg!(it.next()); dbg!(it.next()); dbg!(it.next());
    
    let mut it = s.range(..50).rev();
    dbg!(it.next()); dbg!(it.next()); dbg!(it.next());
}

の出力は50 51 52 49 48 47となる。

ジャッジサーバではCEするが、サンプルでは通したのでたぶん正しいだろう。

fn solve() {
    input!{
        N:usize,
        p:[usize;N],
    }
    let mut set = BTreeSet::new();
    let mut ps=vec![];
    for i in 0..p.len() {
        ps.push((p[i],i+1));
    }
    ps.sort();
    ps.reverse(); // asc

    let mut ans=0;
    for x in ps {
        let p = x.0;
        let i = x.1;
        let mut right=vec![];
        let mut right_it = set.range(i+1..);
        if let Some(x) = right_it.next() {
            right.push(*x);
        }
        if let Some(x) = right_it.next() {
            right.push(*x);
        }
        let mut left=vec![];
        let mut left_it = set.range(..i).rev();
        if let Some(x) = left_it.next() {
            left.push(*x);
        }
        if let Some(x) = left_it.next() {
            left.push(*x);
        }

        while left.len() < 2 {
            left.push(0);
        }
        while right.len() < 2 {
            right.push(N+1);
        }
        left.sort();
        right.sort();

        let ci = (left[1]-left[0])*(right[0]-i) + (right[1]-right[0])*(i-left[1]);
        ans+=p*ci;

        set.insert(i);
    }
    println!("{}",ans);
}

AtCoderのRustバージョンが早急に上がることを願っている。