nekolib/ds/
bicremental_median.rs

1//! 中央値の管理。
2
3use std::collections::BTreeMap;
4use std::fmt::Debug;
5
6/// 中央値の管理。
7///
8/// 多重集合への要素の追加と削除を行いつつ、中央値を管理する。
9///
10/// # Naming
11/// incremental と decremental の双方向の処理を行うので、bidirectional
12/// の気持ちで bicremental とした。記憶が正しければえびちゃんの造語なので、
13/// もっとよい名前があれば変えたい。dynamic は意味が曖昧なのできらい。
14///
15/// # Notes
16/// 集合に $k$ 個追加・削除する操作をサポートできないか？と思ったが、
17/// これだと計算量の保証ができない。$\\{\\{1, 2, \\dots, n\\}\\}$ に対して
18/// $0$ を $n$ 個追加する操作と、$0$ を $n$ 個削除する操作を繰り返すことで、
19/// 簡単に worst $\\Omega(n)$ 時間になってしまう。
20#[derive(Clone, Debug, Eq, PartialEq)]
21pub struct BicrementalMedian<T: Ord + Clone> {
22    lower_len: usize,
23    upper_len: usize,
24    lower: BTreeMap<T, usize>,
25    upper: BTreeMap<T, usize>,
26}
27
28impl<T: Ord + Clone> BicrementalMedian<T> {
29    pub fn new() -> Self {
30        Self {
31            lower_len: 0,
32            upper_len: 0,
33            lower: BTreeMap::new(),
34            upper: BTreeMap::new(),
35        }
36    }
37    pub fn insert(&mut self, x: T) {
38        if self.lower_len == 0 {
39            self.lower.insert(x, 1);
40            self.lower_len += 1;
41        } else if self.lower_len == self.upper_len {
42            // [LLL] X [RRR]
43            if &x <= self.upper.iter().next().unwrap().0 {
44                // [LLXL] [RRR]
45                *self.lower.entry(x).or_insert(0) += 1;
46            } else {
47                // [LLLR] [RRX]
48                self.rotate_to_lower();
49                *self.upper.entry(x).or_insert(0) += 1;
50            }
51            self.lower_len += 1;
52        } else {
53            // [LLL] X [RR]
54            if self.lower.iter().next_back().unwrap().0 < &x {
55                // [LLL] [RXR]
56                *self.upper.entry(x).or_insert(0) += 1;
57            } else {
58                // [XLL] [LRR]
59                self.rotate_to_upper();
60                *self.lower.entry(x).or_insert(0) += 1;
61            }
62            self.upper_len += 1;
63        }
64    }
65    pub fn remove(&mut self, x: &T) -> bool {
66        if self.lower_len == 0 {
67            false
68        } else if self.lower_len == self.upper_len {
69            // [LLL] [RRR]
70            if self.upper.contains_key(x) {
71                // [LLL] [RR]
72                self.remove_from_upper(x, false);
73                return true;
74            }
75            if self.lower.contains_key(x) {
76                // [LLR] [RR]
77                self.remove_from_lower(x, true);
78                return true;
79            }
80            false
81        } else {
82            // [LLL] [RR]
83            if self.lower.contains_key(x) {
84                // [LL] [RR]
85                self.remove_from_lower(x, false);
86                return true;
87            }
88            if self.upper.contains_key(x) {
89                // [LL] [LR]
90                self.remove_from_upper(x, true);
91                return true;
92            }
93            false
94        }
95    }
96    pub fn median(&self) -> Option<&T> {
97        if self.lower_len == 0 {
98            None
99        } else {
100            Some(self.lower.iter().next_back().unwrap().0)
101        }
102    }
103}
104
105impl<T: Ord + Clone> BicrementalMedian<T> {
106    fn rotate_to_lower(&mut self) {
107        let (x, k) =
108            self.upper.iter().next().map(|(x, &k)| (x.clone(), k)).unwrap();
109        if k == 1 {
110            self.upper.remove(&x);
111        } else {
112            *self.upper.get_mut(&x).unwrap() -= 1;
113        }
114        *self.lower.entry(x).or_insert(0) += 1;
115    }
116    fn rotate_to_upper(&mut self) {
117        let (x, k) = self
118            .lower
119            .iter()
120            .next_back()
121            .map(|(x, &k)| (x.clone(), k))
122            .unwrap();
123        if k == 1 {
124            self.lower.remove(&x);
125        } else {
126            *self.lower.get_mut(&x).unwrap() -= 1;
127        }
128        *self.upper.entry(x).or_insert(0) += 1;
129    }
130    fn remove_from_lower(&mut self, x: &T, rotate: bool) {
131        if self.lower[x] == 1 {
132            self.lower.remove(x);
133        } else {
134            *self.lower.get_mut(x).unwrap() -= 1;
135        }
136        if rotate {
137            self.rotate_to_lower();
138            self.upper_len -= 1;
139        } else {
140            self.lower_len -= 1;
141        }
142    }
143    fn remove_from_upper(&mut self, x: &T, rotate: bool) {
144        if self.upper[x] == 1 {
145            self.upper.remove(x);
146        } else {
147            *self.upper.get_mut(x).unwrap() -= 1;
148        }
149        if rotate {
150            self.rotate_to_upper();
151            self.lower_len -= 1;
152        } else {
153            self.upper_len -= 1;
154        }
155    }
156}
157
158#[test]
159fn test_simple() {
160    let n = 32768;
161    let mut f =
162        std::iter::successors(Some(296), |&x| Some((x * 258 + 185) % 397))
163            .map(|x| x & 15);
164    let mut bucket = vec![0; 8];
165    let mut bm = BicrementalMedian::<usize>::new();
166    for _ in 0..n {
167        let x = f.next().unwrap();
168        let (remove, x) = (x & 8 != 0, x & 7);
169        if remove && bucket[x] > 0 {
170            bucket[x] -= 1;
171            bm.remove(&x);
172        } else {
173            bucket[x] += 1;
174            bm.insert(x);
175        }
176        let mut naive = vec![];
177        for i in 0..8 {
178            naive.extend(std::iter::repeat(i).take(bucket[i]));
179        }
180        assert_eq!(bm.median(), naive.get(naive.len().wrapping_sub(1) / 2));
181    }
182}
nekolib/ds/bicremental_median.rs

nekolib/ds/
bicremental_median.rs
