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

• 問題: No.649 ここでちょっとQK！
• ユーザー: Shuntaro Ohno
• 提出日時: 2020-11-25 13:50:07
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 21,428 bytes
• コンパイル時間: 14,394 ms
• コンパイル使用メモリ: 383,940 KB
• 実行使用メモリ: 25,344 KB
• 最終ジャッジ日時: 2024-07-23 19:18:08
• 合計ジャッジ時間: 22,279 ms
• ジャッジサーバーID （参考情報）: judge5 / judge4

ファイルパターン	結果
sample	-- * 4
other	AC * 10 WA * 1 RE * 1 TLE * 1 -- * 19

ソースコード

use std::cmp::max;
use std::cmp::Reverse;

#[derive(Clone, Copy, PartialEq)]
pub enum BinSide {
    Left,
    Right,
    Root,
}

/// Heap que (ascending order).
pub struct HeapQ<T: PartialOrd + Copy> {
    /// Vector of data nodes.
    pub data: Vec<T>
}

impl<T: PartialOrd + Copy> HeapQ<T> {
    /// Create an empty tree.
    pub fn new() -> Self {
        let data = Vec::new();
        Self { data }
    }

    /// Add a new data node.
    pub fn push(&mut self, e: T) {
        let n = self.data.len();
        self.data.push(e);
        self.lookup(n);
    }

    /// Pop a smallest value.
    pub fn pop(&mut self) -> Option<T> {
        let size = self.data.len();
        match size {
            0 => None,
            1 => self.data.pop(),
            n => {
                let tmp = self.data[0];
                self.data[0] = self.data[n-1];
                let _ = self.data.pop();
                self.lookdown(0);
                Some(tmp)
            }
        }
    }

    /// Return the number of the data nodes in heap que.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Return vector of contents (NOTE: not sorted!)
    pub fn to_vec(&self) -> Vec<T> {
        self.data.clone()
    }

    /// Update the tree from the bottom up.
    pub fn lookup(&mut self, n: usize) {
        if n == 0 { return; }
        let p = (n - 1) / 2;
        if self.data[p] > self.data[n] {
            let tmp = self.data[p];
            self.data[p] = self.data[n];
            self.data[n] = tmp;
            self.lookup(p);
        }
    }

    /// Update the tree from the top down.
    pub fn lookdown(&mut self, n: usize) {
        let c1 = 2 * n + 1;
        let c2 = 2 * n + 2;
        let size = self.data.len();
        if c1 >= size { return; }
        if c2 >= size {
            if self.data[n] > self.data[c1] {
                let tmp = self.data[n];
                self.data[n] = self.data[c1];
                self.data[c1] = tmp;
            }
        } else {
            if self.data[c1] < self.data[c2] {
                if self.data[n] > self.data[c1] {
                    let tmp = self.data[n];
                    self.data[n] = self.data[c1];
                    self.data[c1] = tmp;
                    self.lookdown(c1);
                }
            } else {
                if self.data[n] > self.data[c2] {
                    let tmp = self.data[n];
                    self.data[n] = self.data[c2];
                    self.data[c2] = tmp;
                    self.lookdown(c2);
                }
            }
        }
    }
}

/// Heap que (descending order).
pub struct HeapQD<T: PartialOrd + Copy> {
    pub heapq: HeapQ<Reverse<T>>
}

impl<T: PartialOrd + Copy> HeapQD<T> {
    /// Create an empty tree.
    pub fn new() -> Self {
        let heapq = HeapQ::new();
        Self { heapq }
    }

    /// Add a new data node.
    pub fn push(&mut self, e: T) {
        self.heapq.push(Reverse(e));
    }

    /// Pop a smallest value.
    pub fn pop(&mut self) -> Option<T> {
        match self.heapq.pop() {
            Some(Reverse(e)) => Some(e),
            _ => None
        }
    }

    /// Return the number of the data nodes in heap que.
    pub fn len(&self) -> usize {
        self.heapq.len()
    }

    /// Return vector of contents (NOTE: not sorted!)
    pub fn to_vec(&self) -> Vec<T> {
        self.heapq.to_vec().iter().map(|r| r.0).collect::<Vec<T>>()
    }
}

pub struct AVLNode<T: PartialOrd + Copy> {
    pub id: usize,
    pub value: T,
    pub n_dup: usize,
    pub n_lch: usize,
    pub n_rch: usize,

    pub parent: Option<usize>,
    pub lch: Option<usize>,
    pub rch: Option<usize>,

    pub d_lch: usize,
    pub d_rch: usize,
}

pub struct AVLTree<T: PartialOrd + Copy> {
    pub nodes: Vec<Box<AVLNode<T>>>,
    pub next_ids: HeapQ<usize>,
    pub root: Option<usize>
}

impl<T: PartialOrd + Copy> AVLNode<T> {
    pub fn new(id: usize, value: T) -> Self {
        let n_dup = 1;
        let n_lch = 0;
        let n_rch = 0;
        let parent = None;
        let lch = None;
        let rch = None;

        let d_lch = 0;
        let d_rch = 0;

        Self { id, value, n_dup, n_lch, n_rch, parent, lch, rch, d_lch, d_rch }
    }

    pub fn init(&mut self, value: T) {
        self.value = value;
        self.n_dup = 1;
        self.n_lch = 0;
        self.n_rch = 0;
        self.parent = None;
        self.lch = None;
        self.rch = None;

        self.d_lch = 0;
        self.d_rch = 0;
    }
}

impl<T: PartialOrd + Copy> AVLTree<T> {
    pub fn new() -> Self {
        let nodes = Vec::new();
        let mut next_ids = HeapQ::new();
        next_ids.push(0);
        let root = None;
        Self { nodes, next_ids, root }
    }

    pub fn find(&self, e: T) -> Result<usize, (usize, BinSide)> {
        match self.root {
            None => Err((0, BinSide::Root)),
            Some(r) => {
                let mut i = r;
                loop {
                    if self.nodes[i].value == e {
                        return Ok(i);
                    } else if self.nodes[i].value > e {
                        match self.nodes[i].lch {
                            None => return Err((i, BinSide::Left)),
                            Some(j) => i = j,
                        }
                    } else {
                        match self.nodes[i].rch {
                            None => return Err((i, BinSide::Right)),
                            Some(j) => i = j,
                        }
                    }
                }
            }
        }
    }

    pub fn find_count(&self, e: T) -> Result<(usize, usize), (usize, BinSide, usize)> {
        match self.root {
            None => Err((0, BinSide::Root, 0)),
            Some(r) => {
                let mut left = 0;
                let mut i = r;
                loop {
                    left += self.nodes[i].n_lch + self.nodes[i].n_dup;
                    if self.nodes[i].value == e {
                        return Ok((i, left));
                    } else if self.nodes[i].value > e {
                        match self.nodes[i].lch {
                            None => return Err((i, BinSide::Left, left - self.nodes[i].n_dup)),
                            Some(j) => i = j,
                        }
                    } else {
                        match self.nodes[i].rch {
                            None => return Err((i, BinSide::Right, left)),
                            Some(j) => i = j,
                        }
                    }
                }
            }
        }
    }

    pub fn rotate_r(&mut self, u: usize, v: usize) {
        match self.nodes[u].parent {
            None => {
                self.nodes[v].parent = None;
                self.root = Some(v);
            }
            Some(p) => {
                self.nodes[v].parent = Some(p);
                if self.nodes[p].value < self.nodes[v].value {
                    self.nodes[p].rch = Some(v);
                } else {
                    self.nodes[p].lch = Some(v);
                }
            }
        }
        self.nodes[u].parent = Some(v);
        self.nodes[u].lch = self.nodes[v].rch;
        self.nodes[u].n_lch = self.nodes[v].n_rch;
        self.nodes[u].d_lch = self.nodes[v].d_rch;
        self.nodes[v].rch = Some(u);
        self.nodes[v].n_rch = self.nodes[u].n_lch + self.nodes[u].n_rch + self.nodes[u].n_dup;
        self.nodes[v].d_rch = max(self.nodes[u].d_lch, self.nodes[u].d_rch) + 1;
    }

    pub fn rotate_l(&mut self, u: usize, v: usize) {
        match self.nodes[u].parent {
            None => {
                self.nodes[v].parent = None;
                self.root = Some(v);
            }
            Some(p) => {
                self.nodes[v].parent = Some(p);
                if self.nodes[p].value < self.nodes[v].value {
                    self.nodes[p].rch = Some(v);
                } else {
                    self.nodes[p].lch = Some(v);
                }
            }
        }
        self.nodes[u].parent = Some(v);
        self.nodes[u].rch = self.nodes[v].lch;
        self.nodes[u].n_rch = self.nodes[v].n_lch;
        self.nodes[u].d_rch = self.nodes[v].d_lch;
        self.nodes[v].lch = Some(u);
        self.nodes[v].n_lch = self.nodes[u].n_lch + self.nodes[u].n_rch + self.nodes[u].n_dup;
        self.nodes[v].d_lch = max(self.nodes[u].d_lch, self.nodes[u].d_rch) + 1;
    }

    pub fn rotate_rl(&mut self, u: usize, v: usize, w: usize) {
        match self.nodes[u].parent {
            None => {
                self.nodes[w].parent = None;
                self.root = Some(w);
            }
            Some(p) => {
                self.nodes[w].parent = Some(p);
                if self.nodes[p].value < self.nodes[w].value {
                    self.nodes[p].rch = Some(w);
                } else {
                    self.nodes[p].lch = Some(w);
                }
            }
        }
        self.nodes[u].parent = Some(w);
        self.nodes[v].parent = Some(w);
        self.nodes[u].rch = self.nodes[w].lch;
        self.nodes[u].n_rch = self.nodes[w].n_lch;
        self.nodes[u].d_rch = self.nodes[w].d_lch;
        self.nodes[v].lch = self.nodes[w].rch;
        self.nodes[v].n_lch = self.nodes[w].n_rch;
        self.nodes[v].d_lch = self.nodes[w].d_rch;
        self.nodes[w].lch = Some(u);
        self.nodes[w].n_lch = self.nodes[u].n_lch + self.nodes[u].n_rch + self.nodes[u].n_dup;
        self.nodes[w].d_lch = max(self.nodes[u].d_lch, self.nodes[u].d_rch) + 1;
        self.nodes[w].rch = Some(v);
        self.nodes[w].n_rch = self.nodes[v].n_lch + self.nodes[v].n_rch + self.nodes[v].n_dup;
        self.nodes[w].d_rch = max(self.nodes[v].d_lch, self.nodes[v].d_rch) + 1;
    }

    pub fn rotate_lr(&mut self, u: usize, v: usize, w: usize) {
        match self.nodes[u].parent {
            None => {
                self.nodes[w].parent = None;
                self.root = Some(w);
            }
            Some(p) => {
                self.nodes[w].parent = Some(p);
                if self.nodes[p].value < self.nodes[w].value {
                    self.nodes[p].rch = Some(w);
                } else {
                    self.nodes[p].lch = Some(w);
                }
            }
        }
        self.nodes[v].parent = Some(w);
        self.nodes[u].parent = Some(w);
        self.nodes[v].rch = self.nodes[w].lch;
        self.nodes[v].n_rch = self.nodes[w].n_lch;
        self.nodes[v].d_rch = self.nodes[w].d_lch;
        self.nodes[u].lch = self.nodes[w].rch;
        self.nodes[u].n_lch = self.nodes[w].n_rch;
        self.nodes[u].d_lch = self.nodes[w].d_rch;
        self.nodes[w].lch = Some(v);
        self.nodes[w].n_lch = self.nodes[v].n_lch + self.nodes[v].n_rch + self.nodes[v].n_dup;
        self.nodes[w].d_lch = max(self.nodes[v].d_lch, self.nodes[v].d_rch) + 1;
        self.nodes[w].rch = Some(u);
        self.nodes[w].n_rch = self.nodes[u].n_lch + self.nodes[u].n_rch + self.nodes[u].n_dup;
        self.nodes[w].d_rch = max(self.nodes[u].d_lch, self.nodes[u].d_rch) + 1;
    }

    pub fn update_ancestors(&mut self, id: usize) {
        let mut i = id;
        let mut flg = true;
        loop {
            match self.nodes[i].parent {
                None => return,
                Some(j) => {
                    let n = self.nodes[i].n_lch + self.nodes[i].n_rch + self.nodes[i].n_dup;
                    let d = max(self.nodes[i].d_lch, self.nodes[i].d_rch) + 1;
                    if self.nodes[i].value < self.nodes[j].value {
                        self.nodes[j].n_lch = n;
                        self.nodes[j].d_lch = d;
                    } else {
                        self.nodes[j].n_rch = n;
                        self.nodes[j].d_rch = d;
                    }
                    let next = if flg {
                        let bj = self.nodes[j].d_lch as i8 - self.nodes[j].d_rch as i8;
                        let bi = self.nodes[i].d_lch as i8 - self.nodes[i].d_rch as i8;
                        if bj == 2 {
                            flg = false;
                            if bi >= 0 {
                                self.rotate_r(j, i);
                                i
                            } else {
                                let w = self.nodes[i].rch.unwrap();
                                self.rotate_lr(j, i, w);
                                w
                            }
                        } else if bj == -2 {
                            flg = false;
                            if bi <= 0 {
                                self.rotate_l(j, i);
                                i
                            } else {
                                let w = self.nodes[i].lch.unwrap();
                                self.rotate_rl(j, i, w);
                                w
                            }
                        } else {
                            j
                        }
                    } else {  
                        j
                    };
                    i = next;
                }
            }
        }
    }

    pub fn max(&self) -> Option<T> {
        match self.root {
            Some(r) => Some(self.max_from_id(r).0),
            None => None,
        }
    }

    pub fn min(&self) -> Option<T> {
        match self.root {
            Some(r) => Some(self.min_from_id(r).0),
            None => None,
        }
    }

    pub fn max_from_id(&self, id: usize) -> (T, usize) {
        let mut i = id;
        loop {
            match self.nodes[i].rch {
                None => return (self.nodes[i].value, i),
                Some(j) => i = j,
            }
        }
    }

    pub fn min_from_id(&self, id: usize) -> (T, usize) {
        let mut i = id;
        loop {
            match self.nodes[i].lch {
                None => return (self.nodes[i].value, i),
                Some(j) => i = j,
            }
        }
    }

    pub fn kth(&self, k: usize) -> Option<T> {
        let mut i = match self.root {
            None => return None,
            Some(r) => r,
        };
        if k >= self.nodes[i].n_lch + self.nodes[i].n_dup + self.nodes[i].n_rch {
            return None;
        }
        let mut left = 0;
        loop {
            if self.nodes[i].n_lch + left > k {
                i = self.nodes[i].lch.unwrap();
            } else if self.nodes[i].n_lch + self.nodes[i].n_dup + left <= k {
                left += self.nodes[i].n_lch + self.nodes[i].n_dup;
                i = self.nodes[i].rch.unwrap();
            } else {
                return Some(self.nodes[i].value)
            }
        }
    }

    pub fn insert(&mut self, e: T) -> Result<(), usize> {
        match self.find(e) {
            Ok(i) => {
                self.nodes[i].n_dup += 1;
                self.update_ancestors(i);
                Err(self.nodes[i].n_dup)
            }
            Err((i, side)) => {
                let id = self.next_ids.pop().unwrap();
                if self.next_ids.len() == 0 {
                    self.next_ids.push(id + 1);
                    self.nodes.push(Box::new(AVLNode::new(id, e)));
                } else {
                    self.nodes[id].init(e);
                }
                if side == BinSide::Root {
                    self.root = Some(id);
                } else {
                    match side {
                        BinSide::Left => self.nodes[i].lch = Some(id),
                        _ => self.nodes[i].rch = Some(id),
                    };
                    self.nodes[id].parent = Some(i);
                }
                self.update_ancestors(id);
                Ok(())
            }
        }
    }

    pub fn remove(&mut self, e: T) -> Result<T, ()> {
        match self.remove_n(e, 1) {
            Ok((ele, _)) => Ok(ele),
            Err(_) => Err(())
        }
    }

    pub fn remove_all(&mut self, e: T) -> Result<(T, usize), (T, usize)> {
        match self.remove_n(e, std::usize::MAX) {
            Err((ele, 0)) => Err((ele, 0)),
            Err((ele, n)) => Ok((ele, n)),
            Ok(t) => Ok(t),
        }
    }

    pub fn remove_n(&mut self, e: T, n: usize) -> Result<(T, usize), (T, usize)> {
        let i = match self.find(e) {
            Err(_) => return Err((e, 0)),
            Ok(i) => i,
        };
        let n_dup = self.nodes[i].n_dup;
        let res = if n_dup < n {
            self.nodes[i].n_dup = 0;
            Err((e, n_dup))
        } else {
            self.nodes[i].n_dup -= n;
            Ok((e, n_dup))
        };
        if self.nodes[i].n_dup == 0 {
            match (self.nodes[i].lch, self.nodes[i].rch) {
                (Some(lj), Some(_)) => {
                    let (v, j) = self.max_from_id(lj);
                    let pj = self.nodes[j].parent.unwrap();
                    if self.nodes[pj].value > self.nodes[j].value {
                        self.nodes[pj].lch = self.nodes[j].lch;
                        match self.nodes[j].lch {
                            None => {
                                self.nodes[pj].n_lch = 0;
                                self.nodes[pj].d_lch = 0;
                                self.update_ancestors(pj);
                            }
                            Some(lk) => {
                                self.nodes[lk].parent = Some(pj);
                                self.update_ancestors(lk);
                            }
                        }
                    } else {
                        self.nodes[pj].rch = self.nodes[j].lch;
                        match self.nodes[j].lch {
                            None => {
                                self.nodes[pj].n_rch = 0;
                                self.nodes[pj].d_rch = 0;
                                self.update_ancestors(pj);
                            }
                            Some(lk) => {
                                self.nodes[lk].parent = Some(pj);
                                self.update_ancestors(lk);
                            }
                        }
                    }
                    self.nodes[i].value = v;
                    self.nodes[i].n_dup = self.nodes[j].n_dup;
                    self.next_ids.push(j);
                }
                (Some(lj), _) => {
                    self.nodes[lj].parent = self.nodes[i].parent;
                    match self.nodes[i].parent {
                        None => self.root = Some(lj),
                        Some(j) => {
                            if self.nodes[lj].value < self.nodes[j].value {
                                self.nodes[j].lch = Some(lj);
                            } else {
                                self.nodes[j].rch = Some(lj)
                            }
                            self.update_ancestors(j);
                        }
                    };
                    self.next_ids.push(i);
                }
                (_, Some(rj)) => {
                    self.nodes[rj].parent = self.nodes[i].parent;
                    match self.nodes[i].parent {
                        None => self.root = Some(rj),
                        Some(j) => {
                            if self.nodes[rj].value < self.nodes[j].value {
                                self.nodes[j].lch = Some(rj);
                            } else {
                                self.nodes[j].rch = Some(rj)
                            }
                            self.update_ancestors(j);
                        }
                    };
                    self.next_ids.push(i);
                }
                (_, _) => {
                    match self.nodes[i].parent {
                        None => self.root = None,
                        Some(j) => {
                            if self.nodes[i].value < self.nodes[j].value {
                                self.nodes[j].lch = None;
                                self.nodes[j].n_lch = 0;
                                self.nodes[j].d_lch = 0;
                            } else {
                                self.nodes[j].rch = None;
                                self.nodes[j].n_rch = 0;
                                self.nodes[j].d_rch = 0;
                            }
                            self.update_ancestors(j);
                        }
                    };
                    self.next_ids.push(i);
                }
            }
        } else {
            self.update_ancestors(i);
        }
        res
    }
}

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

pub fn readn<T: std::str::FromStr>(delimiter: &str) -> Vec<T> {
    let s = read1::<String>();
    s.split(delimiter).map(|e| e.parse().ok().unwrap()).collect::<Vec<T>>()
}

fn main() {
    let qk = readn::<usize>(" ");
    let q = qk[0];
    let k = qk[1];
    let mut avl = AVLTree::new();
    for _ in 0..q {
        let tx = readn::<usize>(" ");
        let t = tx[0];
        if t == 1 {
            let _ = avl.insert(tx[1]);
        } else {
            let i = match avl.kth(k-1) {
                None => {
                    println!("-1");
                    continue;
                }
                Some(e) => e,
            };
            println!("{}", i);
            let _ = avl.remove(i);
        }
    }
}