Rustコトハジメ

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

【ABC123-D Cake 123】上位K個系の教育的良問

競プロあるあるテクですが、私もはじめに見た時はビビったので、教育的良問として紹介します。

D - Cake 123

問題

A1...AX, B1...BY, C1...CZからそれぞれ1つずつ値を選び、AxBxCについて上位K個を列挙します。

ただし、X,Y,Z=103ですが、Kは3000程度です。

方針

上位1個なら最大値をとってきて終わりです。

しかし上位K個となると、ナイーブに考えると全部計算してからソートして上位K個ですが、間に合いません。

そこで、AxBxCの上位K個ということは、各要素が上位K個に入ってないと明らかにだめでしょ(例えば、AについてK+1個目はとりようがない)ということに気づき、それを発展させて、AxBについて上位K個とCについて上位K個のミックスでええやんということに気づくと、まずはAxBについて全部列挙してから上位K個をとるという方針にたどりつけます。

実装

fn solve() {
    input!{
        X:usize,Y:usize,Z:usize,K:usize,
        XX:[usize;X],
        YY:[usize;Y],
        ZZ:[usize;Z],
    }
    let mut xy_sum = vec![];
    for &x in &XX {
        for &y in &YY {
            xy_sum.push(x+y);
        }
    }
    xy_sum.sort();
    xy_sum.reverse();
 
    let mut xyz_sum = vec![];
    for xy in xy_sum.split_at(min(3000,xy_sum.len())).0 {
        for &z in &ZZ {
            xyz_sum.push(xy+z);
        }
    }
 
    xyz_sum.sort();
    xyz_sum.reverse();
 
    for i in 0..K {
        println!("{}",xyz_sum[i]);
    }
}

考察

半分全列挙でも応用されるあるあるテクニックです。例えば、N=40を20ずつにわけて、それぞれについて網羅的に探索をして、上位K個ずつで合体させるという発想です。