bouzuya.hatenablog.com

ぼうずやのにっき

ABC057 の A, B, C, D を解いた / Iterator::position の複数回呼び出し

ABC057 : AtCoder Beginner Contest 057 の A, B, C, D を解いた。

use proconio::input;

fn choose(n: usize, r: usize) -> usize {
    if n < r {
        return 0;
    }
    if r == 0 {
        return 1;
    }
    let mut m = 1;
    for i in 0..r {
        m *= n - i;
        m /= i + 1;
    }
    m
}

fn main() {
    input! {
        n: usize,
        a: usize,
        b: usize,
        mut v: [usize; n],
    };

    v.sort();
    v.reverse();

    let avg = v.iter().take(a).sum::<usize>() as f64 / a as f64;

    let mut count_eqmin = 0_usize;
    let mut count_fixed = 0_usize;
    for (i, v_i) in v.iter().copied().enumerate() {
        if v_i == v[a - 1] {
            count_eqmin += 1;
            if i < a {
                count_fixed += 1;
            }
        }
    }

    let count = if v[0] == v[a - 1] {
        (a..=b).map(|x| choose(count_eqmin, x)).sum::<usize>()
    } else {
        choose(count_eqmin, count_fixed)
    };

    println!("{}", avg);
    println!("{}", count);
}

Iterator::position は &mut self で iter を進める。複数回呼び出すと前回の最後に見つかった位置の次の位置から探索できることを知った。

fn main() {
    let v = vec![1, 1, 1];
    let mut iter = v.iter();
    assert_eq!(iter.position(|x| x == &1), Some(0)); // next が v[0] を指す iter から探して 0 の位置
    assert_eq!(iter.position(|x| x == &1), Some(0)); // next が v[1] を指す iter から探して 0 の位置
    assert_eq!(iter.position(|x| x == &1), Some(0)); // next が v[2] を指す iter から探して 0 の位置
    assert_eq!(iter.position(|x| x == &1), None);
}

今日のコミット。