Skip to main content

library_checker/data_structure/
ordered_set.rs

1use competitive::prelude::*;
2use competitive::{
3    algorithm::SliceSortExt,
4    data_structure::{DaryPrefixSumTreeU32, StaticSearch},
5};
6
7#[verify::library_checker("ordered_set")]
8pub fn ordered_set(reader: impl Read, writer: impl Write) {
9    prepare_io!(buffered; reader, writer);
10    sc!(n, q, a: [u32; n], queries: [(u8, u32); q]);
11    let mut values: Vec<_> = a
12        .iter()
13        .copied()
14        .chain(queries.iter().filter_map(|&(t, x)| (t <= 1).then_some(x)))
15        .collect();
16    values.radix_sort_by_key(|&x| x);
17    values.dedup();
18    let search = StaticSearch::from_sorted(&values);
19    let endpoints: Vec<_> = a
20        .into_iter()
21        .chain(
22            queries
23                .iter()
24                .map(|&(t, x)| if t == 3 || t == 4 { x + 1 } else { x }),
25        )
26        .collect();
27    let mut positions = vec![0; endpoints.len()];
28    search.lower_bound_batch(&endpoints, &mut positions);
29    let mut counts = vec![0; values.len()];
30    for &k in &positions[..n] {
31        counts[k] = 1;
32    }
33    let mut seg = DaryPrefixSumTreeU32::from_slice(&counts);
34    for ((t, x), &k) in queries.into_iter().zip(&positions[n..]) {
35        match t {
36            0 => seg.set(k, 1),
37            1 => seg.set(k, 0),
38            2 => {
39                let k = seg.partition_point_acc(x - 1);
40                pp!(values.get(k).map_or(-1, |&x| x as i64));
41            }
42            3 => {
43                pp!(seg.accumulate0(k));
44            }
45            4 => {
46                let count = seg.accumulate0(k);
47                pp!(if count == 0 {
48                    -1
49                } else {
50                    values[seg.partition_point_acc(count - 1)] as i64
51                });
52            }
53            5 => {
54                let count = seg.accumulate0(k);
55                let k = seg.partition_point_acc(count);
56                pp!(values.get(k).map_or(-1, |&x| x as i64));
57            }
58            _ => unreachable!(),
59        }
60    }
61}