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 { /// Vector of data nodes. pub data: Vec } impl HeapQ { /// 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 { 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 { 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 { pub heapq: HeapQ> } impl HeapQD { /// 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 { 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 { self.heapq.to_vec().iter().map(|r| r.0).collect::>() } } pub struct AVLNode { pub id: usize, pub value: T, pub n_dup: usize, pub n_lch: usize, pub n_rch: usize, pub parent: Option, pub lch: Option, pub rch: Option, pub d_lch: usize, pub d_rch: usize, } pub struct AVLTree { pub nodes: Vec>>, pub next_ids: HeapQ, pub root: Option } impl AVLNode { 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 AVLTree { 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 { 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; match self.nodes[v].rch { None => {} Some(t) => self.nodes[t].parent = Some(u), } 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; match self.nodes[v].lch { None => {} Some(t) => self.nodes[t].parent = Some(u), } 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; match self.nodes[w].lch { None => {} Some(t) => self.nodes[t].parent = Some(u), } 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; match self.nodes[w].rch { None => {} Some(t) => self.nodes[t].parent = Some(v), } 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 { match self.root { Some(r) => Some(self.max_from_id(r).0), None => None, } } pub fn min(&self) -> Option { 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 { 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 { 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(lj); } }; 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(rj); } }; 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 { let mut s = String::new(); std::io::stdin().read_line(&mut s).ok(); s.trim().parse().ok().unwrap() } pub fn readn(delimiter: &str) -> Vec { let s = read1::(); s.split(delimiter).map(|e| e.parse().ok().unwrap()).collect::>() } fn main() { let qk = readn::(" "); let q = qk[0]; let k = qk[1]; let mut avl = AVLTree::new(); for _ in 0..q { let tx = readn::(" "); 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); } } }