AtCoder Beginner Contest 306 E

https://atcoder.jp/contests/abc306/tasks/abc306_e

平衡木を実装できれば簡単なのですが、そうでなくてもK番目前後で分けて、出し入れが管理できていればよいです。そのために、2つBtreeMapを用意して、クエリーのたびにそれらをアップデートします。

// Best Performances
#![allow(non_snake_case)]

use std::collections::BTreeMap;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, usize, usize) {
    let v = read_vec();
    let N = v[0];
    let K = v[1];
    let Q = v[2];
    (N, K, Q)
}

fn read_query() -> (usize, i64) {
    let v: Vec<usize> = read_vec();
    let X = v[0] - 1;
    let Y = v[1] as i64;
    (X, Y)
}

fn insert(n: i64, dic: &mut BTreeMap<i64, i64>) {
    let e = dic.entry(n).or_insert(0);
    *e += 1
}

fn remove(n: &i64, dic: &mut BTreeMap<i64, i64>) {
    if let Some(value) = dic.get_mut(n) {
        *value -= 1;
        if *value == 0 {
            dic.remove(n);
        }
    }
}

fn f(N: usize, K: usize, Q: usize) {
    let mut a: Vec<i64> = (0..N).map(|_| 0).collect();
    let mut s: i64 = 0;
    let mut dic1: BTreeMap<i64, i64> = BTreeMap::new();
    dic1.insert(0, K as i64);
    let mut dic2: BTreeMap<i64, i64> = BTreeMap::new();
    dic2.insert(0, (N - K) as i64);
    for _ in 0..Q {
        let (X, Y) = read_query();
        let (&min1, &_) = dic1.iter().next().unwrap();
        let (&max2, &_) = dic2.iter().rev().next().unwrap();
        if a[X] <= max2 {
            remove(&a[X], &mut dic2);
            if Y <= min1 {
                insert(Y, &mut dic2);
            }
            else {
                remove(&min1, &mut dic1);
                insert(Y, &mut dic1);
                insert(min1, &mut dic2);
                s += Y - min1
            }
        }
        else {
            remove(&a[X], &mut dic1);
            if Y < max2 {
                remove(&max2, &mut dic2);
                insert(Y, &mut dic2);
                insert(max2, &mut dic1);
                s += max2 - a[X]
            }
            else {
                insert(Y, &mut dic1);
                s += Y - a[X]
            }
        }
        a[X] = Y;
        println!("{}", s)
    }
}


//////////////////// main ////////////////////

fn main() {
    let (N, K, Q) = read_input();
    f(N, K, Q)
}