#387398 (Rust) No.649 ここでちょっとQK！

• 問題: No.649 ここでちょっとQK！
• ユーザー: akakimidori
• 提出日時: 2019-10-06 15:53:40
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 128 ms / 3,000 ms
• コード長: 7,859 bytes
• コンパイル時間: 13,467 ms
• コンパイル使用メモリ: 385,192 KB
• 実行使用メモリ: 15,488 KB
• 最終ジャッジ日時: 2024-10-10 04:14:24
• 合計ジャッジ時間: 17,915 ms
• ジャッジサーバーID （参考情報）: judge5 / judge1

ファイルパターン	結果
sample	AC * 4
other	AC * 32

ソースコード

#[allow(dead_code)]
mod tree {
    use std;
    pub struct WBT<T> {
        root: Link<T>,
    }
    impl<T: Ord> WBT<T> {
        pub fn new() -> WBT<T> {
            WBT {
                root: None,
            }
        }
        pub fn insert(&mut self, val: T) {
            let r = self.root.take();
            self.root = Node::insert(r, val);
        }
        pub fn remove(&mut self, val: &T) {
            let r = self.root.take();
            self.root = Node::remove(r, val);
        }
        pub fn search_at(&self, k: usize) -> Option<&T> {
            Node::search_at(&self.root, k)
        }
        pub fn search_by(&self, val: &T) -> usize {
            Node::search_by(&self.root, val)
        }
    }
    type NonNull<T> = Box<Node<T>>;
    type Link<T> = Option<NonNull<T>>;
    struct Node<T> {
        val: T,
        size: usize,
        left: Link<T>,
        right: Link<T>,
    }
    #[derive(PartialEq, Eq)]
    enum Bias {
        N, L, R,
    }
    impl<T: Ord> Node<T> {
        fn get_size(t: &Link<T>) -> usize {
            match t.as_ref() {
                None => 0,
                Some(t) => t.size,
            }
        }
        fn update(&mut self) {
            self.size = 1 + Node::get_size(&self.left) + Node::get_size(&self.right);
        }
        fn get_bias(t: &NonNull<T>, p: usize) -> Bias {
            let l_size = Node::get_size(&t.left) + 1;
            let r_size = Node::get_size(&t.right) + 1;
            if p * l_size >= r_size && l_size <= p * r_size {
                Bias::N
            } else if l_size > p * r_size {
                Bias::L
            } else {
                Bias::R
            }
        }
        fn get_bias_lr(l: &Link<T>, r: &Link<T>, p: usize) -> Bias {
            let l_size = Node::get_size(l) + 1;
            let r_size = Node::get_size(r) + 1;
            if p * l_size >= r_size && l_size <= p * r_size {
                Bias::N
            } else if l_size > p * r_size {
                Bias::L
            } else {
                Bias::R
            }
        }
        fn left_rotate(mut u: NonNull<T>) -> NonNull<T> {
            let mut v = u.right.take().expect("left_rotate: input right child is None");
            u.right = v.left.take();
            u.update();
            v.left = Some(u);
            v.update();
            v
        }
        fn right_rotate(mut u: NonNull<T>) -> NonNull<T> {
            let mut v = u.left.take().expect("right_rotate: input left child is None");
            u.left = v.right.take();
            u.update();
            v.right = Some(u);
            v.update();
            v
        }
        fn balance(mut t: NonNull<T>) -> NonNull<T> {
            match Node::get_bias(&t, 3) {
                Bias::N => t,
                Bias::L => {
                    let l = t.left.take().unwrap();
                    t.left = Some(if Node::get_bias_lr(&l.left, &l.right, 2) == Bias::R {
                        Node::left_rotate(l)
                    } else {
                        l
                    });
                    Node::right_rotate(t)
                },
                Bias::R => {
                    let r = t.right.take().unwrap();
                    t.right = Some(if Node::get_bias_lr(&r.left, &r.right, 2) == Bias::L {
                        Node::right_rotate(r)
                    } else {
                        r
                    });
                    Node::left_rotate(t)
                },
            }
        }
        fn merge(l: Link<T>, r: Link<T>) -> Link<T> {
            if l.is_none() {
                return r;
            }
            if r.is_none() {
                return l;
            }
            Some(match Node::get_bias_lr(&l, &r, 1) {
                Bias::L => {
                    let mut l = l.unwrap();
                    let lr = l.right.take();
                    l.right = Node::merge(lr, r);
                    l.update();
                    Node::balance(l)
                },
                _ => {
                    let mut r = r.unwrap();
                    let rl = r.left.take();
                    r.left = Node::merge(l, rl);
                    r.update();
                    Node::balance(r)
                },
            })
        }
        fn insert(t: Link<T>, val: T) -> Link<T> {
            if t.is_none() {
                return Some(Box::new(Node {
                    val: val,
                    size: 1,
                    left: None,
                    right: None,
                }));
            }
            let mut t = t.unwrap();
            if val < t.val {
                let l = t.left.take();
                t.left = Node::insert(l, val);
            } else {
                let r = t.right.take();
                t.right = Node::insert(r, val);
            }
            t.update();
            Some(Node::balance(t))
        }
        fn remove(t: Link<T>, val: &T) -> Link<T> {
            if t.is_none() {
                return t;
            }
            let mut t = t.unwrap();
            match t.val.cmp(val) {
                std::cmp::Ordering::Equal => {
                    let l = t.left.take();
                    let r = t.right.take();
                    Node::merge(l, r)
                },
                std::cmp::Ordering::Greater => {
                    let l = t.left.take();
                    t.left = Node::remove(l, val);
                    t.update();
                    Some(Node::balance(t))
                },
                std::cmp::Ordering::Less => {
                    let r = t.right.take();
                    t.right = Node::remove(r, val);
                    t.update();
                    Some(Node::balance(t))
                },
            }
        }
        fn search_by(t: &Link<T>, val: &T) -> usize {
            if t.is_none() {
                return 0;
            }
            let t = t.as_ref().unwrap();
            match t.val.cmp(val) {
                std::cmp::Ordering::Equal => {
                    let l_size = Node::get_size(&t.left);
                    l_size
                },
                std::cmp::Ordering::Greater => {
                    Node::search_by(&t.left, val)
                },
                std::cmp::Ordering::Less => {
                    let l_size = Node::get_size(&t.left);
                    l_size + 1 + Node::search_by(&t.right, val)
                },
            }
        }
        fn search_at(t: &Link<T>, k: usize) -> Option<&T> {
            if k >= Node::get_size(t) {
                return None;
            }
            let t = t.as_ref().unwrap();
            let l_size = Node::get_size(&t.left);
            if k == l_size {
                Some(&t.val)
            } else if k < l_size {
                Node::search_at(&t.left, k)
            } else {
                Node::search_at(&t.right, k - l_size - 1)
            }
        }
    }
}

use std::io::{Read, Write};

fn run() {
    let out = std::io::stdout();
    let mut out = std::io::BufWriter::new(out.lock());
    let mut s = String::new();
    std::io::stdin().read_to_string(&mut s).unwrap();
    let mut it = s.trim().split_whitespace();
    let q: usize = it.next().unwrap().parse().unwrap();
    let k: usize = it.next().unwrap().parse().unwrap();
    let mut set = tree::WBT::<u64>::new();
    for _ in 0..q {
        let op: u8 = it.next().unwrap().parse().unwrap();
        if op == 1 {
            let v: u64 = it.next().unwrap().parse().unwrap();
            set.insert(v);
        } else {
            match set.search_at(k - 1) {
                None => writeln!(out, "-1").unwrap(),
                Some(&v) => {
                    set.remove(&v);
                    writeln!(out, "{}", v).unwrap();
                },
            }
        }
    }
}

fn main(){
    run();
}