AtCoder Beginner Contest 330 E

https://atcoder.jp/contests/abc330/tasks/abc330_e

連続する範囲を葉にする木を作ればいいですが、実装が大変です。ですが、範囲にOrderingの各Traitを実装すればBTreeSetを使えます。範囲が重なれば等しい、そうでなければ大小ということにします。

// Counting Ls
#![allow(non_snake_case)]

use std::collections::{BTreeSet, BTreeMap};
use std::cmp::Ordering;


//////////////////// 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()
}


//////////////////// Range ////////////////////

#[derive(Clone)]
struct Range(i32, i32);

impl PartialEq for Range {
    fn eq(&self, other: &Self) -> bool {
        !(self.1 <= other.0 || other.1 <= self.0)
    }
}

impl Eq for Range { }

impl PartialOrd for Range {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Range {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.1 <= other.0 {
            Ordering::Less
        }
        else if other.1 <= self.0 {
            Ordering::Greater
        }
        else {
            Ordering::Equal
        }
    }
}


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

fn read_input() -> (usize, Vec<i32>) {
    let v = read_vec();
    let Q = v[1];
    let A = read_vec();
    (Q, A)
}

fn read_query() -> (usize, i32) {
    let v = read_vec();
    let (i, x) = (v[0] as usize, v[1]);
    (i, x)
}

fn count(A: &Vec<i32>) -> BTreeMap<i32, usize> {
    let mut counter: BTreeMap<i32, usize> = BTreeMap::new();
    for &e in A.iter() {
        let entry = counter.entry(e).or_insert(0);
        *entry += 1
    }
    counter
}

fn make_continuous(A: &Vec<i32>) -> Vec<Range> {
    let mut ranges: Vec<Range> = vec![];
    let mut first = A[0];
    let mut last = first + 1;
    for i in 1..A.len() {
        if last == A[i] {
            last += 1
        }
        else {
            ranges.push(Range(first, last));
            first = A[i];
            last = first + 1
        }
    }
    ranges.push(Range(first, last));
    ranges
}

fn make_tree(counter: &BTreeMap<i32, usize>) -> BTreeSet<Range> {
    let nums: Vec<i32> = counter.iter().map(|(e, _)| *e).collect();
    let ranges = make_continuous(&nums);
    let mut tree: BTreeSet<Range> = BTreeSet::new();
    for range in ranges.into_iter() {
        tree.insert(range);
    }
    tree
}

fn mex(set: &BTreeSet<Range>) -> i32 {
    match set.first() {
        Some(rng) => if rng.0 == 0 { rng.1 } else { 0 },
        None      => 0
    }
}

fn insert(x: i32, tree: &mut BTreeSet<Range>) {
    let op_rng1_ = tree.get(&Range(x - 1, x));
    let op_rng2_ = tree.get(&Range(x + 1, x + 2));
    let (op_rng1, op_rng2) = (op_rng1_.cloned(), op_rng2_.cloned());
    if op_rng1.is_none() && op_rng2.is_none() {
        tree.insert(Range(x, x + 1));
    }
    else if op_rng1.is_none() {
        let rng2 = op_rng2.unwrap();
        tree.remove(&rng2);
        tree.insert(Range(x, rng2.1));
    }
    else if op_rng2.is_none() {
        let rng1 = op_rng1.unwrap();
        tree.remove(&rng1);
        tree.insert(Range(rng1.0, x + 1));
    }
    else {
        let rng1 = op_rng1.unwrap();
        let rng2 = op_rng2.unwrap();
        tree.remove(&rng1);
        tree.remove(&rng2);
        tree.insert(Range(rng1.0, rng2.1));
    }
}

fn remove(y: i32, tree: &mut BTreeSet<Range>) {
    let rng_ = tree.get(&Range(y, y + 1)).unwrap();
    let rng = rng_.clone();
    tree.remove(&rng);
    if y == rng.0 {
        if y != rng.1 - 1 {
            tree.insert(Range(y + 1, rng.1));
        }
    }
    else if y == rng.1 - 1 {
        tree.insert(Range(rng.0, rng.1 - 1));
    }
    else {
        tree.insert(Range(rng.0, y));
        tree.insert(Range(y + 1, rng.1));
    }
}

fn F(Q: usize, mut A: Vec<i32>) {
    let mut counter = count(&A);
    let mut tree: BTreeSet<Range> = make_tree(&counter);
    for _ in 0..Q {
        let (i, x) = read_query();
        let y = A[i-1];
        if x != y {
            if !counter.contains_key(&x) {
                insert(x, &mut tree)
            }
            let e1 = counter.entry(x).or_insert(0);
            *e1 += 1;
            
            let e2 = counter.entry(y).or_insert(0);
            *e2 -= 1;
            if *e2 == 0 {
                remove(y, &mut tree);
                counter.remove(&y);
            }
            A[i-1] = x
        }
        println!("{}", mex(&tree))
    }
}

fn main() {
    let (Q, A) = read_input();
    F(Q, A)
}