nekolib/ds/
interval_set.rs

1//! 区間の集合。
2
3use std::cmp::Ordering::{self, *};
4use std::collections::BTreeSet;
5use std::fmt::Debug;
6use std::ops::{
7    Bound::{self, *},
8    RangeBounds,
9};
10
11/// 区間の集合。
12///
13/// # Notes
14/// 整数のとき、Excluded(x) と Included(x-1) などの扱いに注意。
15/// あくまで実数の区間であるかのように扱われる。
16#[derive(Clone, Debug, Eq)]
17struct Interval<T: Ord>(Bound<T>, Bound<T>);
18
19impl<T: Clone + Ord> From<(Bound<&T>, Bound<&T>)> for Interval<T> {
20    fn from(r: (Bound<&T>, Bound<&T>)) -> Interval<T> {
21        let s = match r.0 {
22            Included(lo) => Included(lo.clone()),
23            Excluded(lo) => Excluded(lo.clone()),
24            Unbounded => Unbounded,
25        };
26        let e = match r.1 {
27            Included(hi) => Included(hi.clone()),
28            Excluded(hi) => Excluded(hi.clone()),
29            Unbounded => Unbounded,
30        };
31        Interval(s, e)
32    }
33}
34
35impl<T: Ord> Interval<T> {
36    fn is_empty(&self) -> bool {
37        match (&self.0, &self.1) {
38            (Unbounded, _) | (_, Unbounded) => false,
39            (Included(lo), Included(hi)) => !(lo <= hi),
40            (Included(lo), Excluded(hi))
41            | (Excluded(lo), Included(hi))
42            | (Excluded(lo), Excluded(hi)) => !(lo < hi),
43        }
44    }
45    fn is_superset(&self, other: &Self) -> bool {
46        if other.is_empty() {
47            return true;
48        }
49        if self.is_empty() {
50            return false;
51        }
52
53        // self.0 <= other.0
54        match (&self.0, &other.0) {
55            (Unbounded, _) => {}
56            (_, Unbounded) => return false,
57            (Excluded(lhs), Included(rhs)) if lhs == rhs => return false,
58            (Included(lhs), Included(rhs))
59            | (Included(lhs), Excluded(rhs))
60            | (Excluded(lhs), Included(rhs))
61            | (Excluded(lhs), Excluded(rhs))
62                if lhs > rhs =>
63            {
64                return false
65            }
66            _ => {}
67        }
68
69        // other.1 <= self.1
70        match (&self.1, &other.1) {
71            (Unbounded, _) => true,
72            (_, Unbounded) => false,
73            (Excluded(lhs), Included(rhs)) => lhs > rhs,
74            (Included(lhs), Included(rhs))
75            | (Included(lhs), Excluded(rhs))
76            | (Excluded(lhs), Excluded(rhs)) => lhs >= rhs,
77        }
78    }
79    fn touches(&self, other: &Self) -> bool {
80        let (left, right) = match self.cmp(&other) {
81            Less => (&self, &other),
82            Equal => return true,
83            Greater => (&other, &self),
84        };
85        match (&left.1, &right.0) {
86            (Unbounded, _) | (_, Unbounded) => true,
87            (Included(le), Included(rs))
88            | (Included(le), Excluded(rs))
89            | (Excluded(le), Included(rs)) => rs <= le,
90            (Excluded(le), Excluded(rs)) => rs < le,
91        }
92    }
93}
94
95fn toggle_bound<T: Ord>(b: Bound<T>) -> Bound<T> {
96    match b {
97        Included(x) => Excluded(x),
98        Excluded(x) => Included(x),
99        Unbounded => Unbounded,
100    }
101}
102
103impl<T: Ord> Ord for Interval<T> {
104    fn cmp(&self, other: &Self) -> Ordering {
105        if self.is_empty() && other.is_empty() {
106            return Equal;
107        }
108        if self.0 != other.0 {
109            return match (&self.0, &other.0) {
110                (Unbounded, _) => Less,
111                (_, Unbounded) => Greater,
112                (Included(lhs), Excluded(rhs)) if lhs == rhs => Less,
113                (Excluded(lhs), Included(rhs)) if lhs == rhs => Greater,
114                (Included(lhs), Included(rhs))
115                | (Included(lhs), Excluded(rhs))
116                | (Excluded(lhs), Included(rhs))
117                | (Excluded(lhs), Excluded(rhs)) => lhs.cmp(rhs),
118            };
119        }
120        if self.1 != other.1 {
121            return match (&self.1, &other.1) {
122                (_, Unbounded) => Less,
123                (Unbounded, _) => Greater,
124                (Excluded(lhs), Included(rhs)) if lhs == rhs => Less,
125                (Included(lhs), Excluded(rhs)) if lhs == rhs => Greater,
126                (Included(lhs), Included(rhs))
127                | (Included(lhs), Excluded(rhs))
128                | (Excluded(lhs), Included(rhs))
129                | (Excluded(lhs), Excluded(rhs)) => lhs.cmp(rhs),
130            };
131        }
132        Equal
133    }
134}
135
136impl<T: Ord> PartialOrd for Interval<T> {
137    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
138        Some(self.cmp(other))
139    }
140}
141
142impl<T: Ord> PartialEq for Interval<T> {
143    fn eq(&self, other: &Self) -> bool { self.cmp(other) == Equal }
144}
145
146/// 区間の集合。
147///
148/// 区間の追加・削除を行うことができる。
149#[derive(Clone, Debug, Eq, PartialEq)]
150pub struct IntervalSet<T: Ord> {
151    buf: BTreeSet<Interval<T>>,
152}
153
154impl<T: Clone + Debug + Ord> IntervalSet<T> {
155    /// 空集合で初期化する。
156    pub fn new() -> Self { Self { buf: BTreeSet::new() } }
157
158    /// 集合が空であれば `true` を返す。
159    pub fn is_empty(&self) -> bool { self.buf.is_empty() }
160
161    /// 区間 `r` を追加する。
162    pub fn insert<R: RangeBounds<T>>(&mut self, r: R) {
163        let mut r: Interval<T> = (r.start_bound(), r.end_bound()).into();
164        if r.is_empty() {
165            return;
166        }
167        self.remove_subset(&r);
168        match self.buf.range(..&r).cloned().next_back() {
169            Some(x) if x.is_superset(&r) => return,
170            Some(x) if x.touches(&r) => {
171                self.buf.remove(&x);
172                r.0 = x.0;
173            }
174            _ => {}
175        }
176        match self.buf.range(&r..).cloned().next() {
177            Some(x) if x.touches(&r) => {
178                self.buf.remove(&x);
179                r.1 = x.1;
180            }
181            _ => {}
182        }
183        self.buf.insert(r);
184    }
185
186    fn insert_if_nonempty(&mut self, x: Interval<T>) {
187        if !x.is_empty() {
188            self.buf.insert(x);
189        }
190    }
191
192    /// 区間 `r` を削除する。
193    pub fn remove<R: RangeBounds<T>>(&mut self, r: R) {
194        let r: Interval<T> = (r.start_bound(), r.end_bound()).into();
195        if r.is_empty() {
196            return;
197        }
198        self.remove_subset(&r);
199        match self.buf.range(..&r).cloned().next_back() {
200            Some(x) if x.is_superset(&r) => {
201                self.buf.remove(&x);
202                let Interval(r0, r1) = r;
203                self.insert_if_nonempty(Interval(x.0, toggle_bound(r0)));
204                self.insert_if_nonempty(Interval(toggle_bound(r1), x.1));
205                return;
206            }
207            Some(mut x) if x.touches(&r) => {
208                self.buf.remove(&x);
209                x.1 = toggle_bound(r.0.clone());
210                self.insert_if_nonempty(x);
211            }
212            _ => {}
213        }
214        match self.buf.range(&r..).cloned().next() {
215            Some(mut x) if x.touches(&r) => {
216                self.buf.remove(&x);
217                x.0 = toggle_bound(r.1);
218                self.insert_if_nonempty(x);
219            }
220            _ => {}
221        }
222    }
223
224    /// 空集合に戻す。
225    pub fn clear(&mut self) { self.buf.clear(); }
226
227    /// `x` 以上の値で、集合中の区間に含まれない最小のものを返す。
228    ///
229    /// # Examples
230    /// ```
231    /// use std::ops::Bound::{Included, Excluded, Unbounded};
232    ///
233    /// use nekolib::ds::IntervalSet;
234    ///
235    /// let mut s = IntervalSet::new();
236    /// s.insert(1..5);
237    /// s.insert(7..=10);
238    /// s.insert(15..);
239    ///
240    /// assert_eq!(s.mex(&0), Included(&0));
241    /// assert_eq!(s.mex(&1), Included(&5));
242    /// assert_eq!(s.mex(&6), Included(&6));
243    /// assert_eq!(s.mex(&7), Excluded(&10));
244    /// assert_eq!(s.mex(&15), Unbounded);
245    /// ```
246    pub fn mex<'a>(&'a self, x: &'a T) -> Bound<&'a T> {
247        if self.buf.is_empty() {
248            return Included(&x);
249        }
250        match self
251            .buf
252            .range(..=Interval(Included(x.clone()), Unbounded))
253            .next_back()
254        {
255            Some(Interval(_, Included(y))) if y < x => Included(x),
256            Some(Interval(_, Included(y))) => Excluded(y),
257            Some(Interval(_, Excluded(y))) if y <= x => Included(x),
258            Some(Interval(_, Excluded(y))) => Included(y),
259            Some(Interval(_, Unbounded)) => Unbounded,
260            None => Included(x),
261        }
262    }
263
264    /// 区間 `r` を含む区間の両端を返す。
265    pub fn covering<R: RangeBounds<T>>(
266        &self,
267        r: &R,
268    ) -> Option<(&Bound<T>, &Bound<T>)> {
269        let r: Interval<T> = (r.start_bound(), r.end_bound()).into();
270        if self.buf.is_empty() {
271            return None;
272        }
273        (if r.is_empty() {
274            self.buf.range(..).next()
275        } else {
276            match self
277                .buf
278                .range(..=&Interval(r.0.clone(), Unbounded))
279                .next_back()
280            {
281                Some(s) if s.is_superset(&r) => Some(s),
282                _ => None,
283            }
284        })
285        .map(|r| (&r.0, &r.1))
286    }
287
288    /// 区間 `r` を含んでいれば `true` を返す。
289    pub fn has_range<R: RangeBounds<T>>(&self, r: &R) -> bool {
290        self.covering(r).is_some()
291    }
292
293    fn remove_subset(&mut self, r: &Interval<T>) {
294        let rem: Vec<Interval<T>> = match r {
295            Interval(Unbounded, Unbounded) => {
296                self.buf.clear();
297                return;
298            }
299            Interval(Included(lo), Unbounded)
300            | Interval(Excluded(lo), Unbounded) => self.buf.range((
301                Included(Interval(Included(lo.clone()), Included(lo.clone()))),
302                Unbounded,
303            )),
304            Interval(Unbounded, Included(hi))
305            | Interval(Unbounded, Excluded(hi)) => self.buf.range((
306                Unbounded,
307                Included(Interval(Included(hi.clone()), Included(hi.clone()))),
308            )),
309            Interval(Included(lo), Included(hi))
310            | Interval(Included(lo), Excluded(hi))
311            | Interval(Excluded(lo), Included(hi))
312            | Interval(Excluded(lo), Excluded(hi)) => self.buf.range((
313                Included(Interval(Included(lo.clone()), Included(lo.clone()))),
314                Included(Interval(Included(hi.clone()), Included(hi.clone()))),
315            )),
316        }
317        .cloned()
318        .collect();
319        for k in rem.into_iter().filter(|x| r.is_superset(x)) {
320            self.buf.remove(&k);
321        }
322    }
323
324    pub fn iter(
325        &self,
326    ) -> impl Iterator<Item = (&Bound<T>, &Bound<T>)> + DoubleEndedIterator + '_
327    {
328        self.buf.iter().map(|x| (&x.0, &x.1))
329    }
330}
nekolib/ds/interval_set.rs

nekolib/ds/
interval_set.rs
