haar_lib/ds/
persistent_stack.rs

1//! 永続スタック
2//!
3//! # Verification
4//! - [ABC273 E - Notebook #37628467](https://atcoder.jp/contests/abc273/submissions/37628467)
5use std::{iter::from_fn, rc::Rc};
6
7#[derive(Debug, Default, Clone)]
8struct Node<T> {
9    value: T,
10    next: Option<Rc<Node<T>>>,
11}
12
13/// 永続スタック
14#[derive(Debug, Default, Clone)]
15pub struct PersistentStack<T> {
16    root: Option<Rc<Node<T>>>,
17}
18
19impl<T> PersistentStack<T> {
20    /// 空の[`PersistentStack`]を生成する。
21    pub fn new() -> Self {
22        Self { root: None }
23    }
24
25    /// 末尾の要素への参照を返す。
26    pub fn peek(&self) -> Option<&T> {
27        self.root.as_ref().map(|x| &x.value)
28    }
29
30    /// 値`value`を末尾に追加した[`PersistentStack`]を返す。
31    pub fn push(&self, value: T) -> Self {
32        Self {
33            root: Some(Rc::new(Node {
34                value,
35                next: self.root.clone(),
36            })),
37        }
38    }
39
40    /// 末尾の要素を削除した[`PersistentStack`]を返す。
41    pub fn pop(&self) -> Option<Self> {
42        self.root.as_ref().map(|root| Self {
43            root: root.next.clone(),
44        })
45    }
46
47    /// スタックの末尾から先頭への要素の参照を返すイテレータを返す。
48    pub fn iter(&self) -> impl Iterator<Item = &T> {
49        let mut root = &self.root;
50
51        from_fn(move || match root {
52            None => None,
53            Some(r) => {
54                root = &r.next;
55                Some(&r.value)
56            }
57        })
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn test() {
67        let a = PersistentStack::<u32>::new();
68        let a = a.push(4).push(5).push(9);
69
70        assert_eq!(a.iter().cloned().collect::<Vec<_>>(), &[9, 5, 4]);
71
72        let b = a.pop().unwrap();
73        assert_eq!(a.iter().cloned().collect::<Vec<_>>(), &[9, 5, 4]);
74        assert_eq!(b.iter().cloned().collect::<Vec<_>>(), &[5, 4]);
75
76        let c = b.push(2);
77        assert_eq!(a.iter().cloned().collect::<Vec<_>>(), &[9, 5, 4]);
78        assert_eq!(b.iter().cloned().collect::<Vec<_>>(), &[5, 4]);
79        assert_eq!(c.iter().cloned().collect::<Vec<_>>(), &[2, 5, 4]);
80    }
81}
haar_lib/ds/persistent_stack.rs

haar_lib/ds/
persistent_stack.rs
