competitive/data_structure/
submask_range_query.rs

1use super::{BitDpExt, Group, Xorshift};
2use std::fmt::{self, Debug};
3
4#[derive(Debug, Clone, Copy)]
5pub struct SubmaskRangeQuery {
6    bit_width: u32,
7    mask: [u32; 3],
8}
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum QueryKind {
12    Get,
13    Update,
14}
15
16impl SubmaskRangeQuery {
17    pub fn new(bit_width: u32) -> Self {
18        let mut rng = Xorshift::default();
19        let mut mask = [0; 3];
20        let mut rem: Vec<_> = (0..bit_width).map(|w| w % 3).collect();
21        rng.shuffle(&mut rem);
22        for (k, r) in rem.into_iter().enumerate() {
23            mask[r as usize] |= 1 << k;
24        }
25        Self { bit_width, mask }
26    }
27
28    pub fn new_with_queries(
29        queries: impl IntoIterator<Item = (QueryKind, u32)> + ExactSizeIterator + Clone,
30    ) -> Self {
31        let bit_width = queries
32            .clone()
33            .into_iter()
34            .map(|(_, m)| 32 - m.leading_zeros())
35            .max()
36            .unwrap_or(0);
37        let mut mask = [0; 3];
38        let mut cost = vec![1u32; queries.len()];
39        for k in 0..bit_width {
40            let mut sum = [0u64; 3];
41            for ((kind, m), &c) in queries.clone().into_iter().zip(&cost) {
42                match kind {
43                    QueryKind::Get => {
44                        let b = m >> k & 1 == 0;
45                        sum[if b { 2 } else { 1 }] += c as u64;
46                    }
47                    QueryKind::Update => {
48                        let b = m >> k & 1 == 0;
49                        sum[if b { 0 } else { 2 }] += c as u64;
50                    }
51                }
52            }
53            let t = (0..3).min_by_key(|&i| sum[i]).unwrap();
54            mask[t] |= 1 << k;
55            for ((kind, m), c) in queries.clone().into_iter().zip(&mut cost) {
56                match kind {
57                    QueryKind::Get => {
58                        let b = m >> k & 1 == 0;
59                        if t == if b { 2 } else { 1 } {
60                            *c <<= 1;
61                        }
62                    }
63                    QueryKind::Update => {
64                        let b = m >> k & 1 == 0;
65                        if t == if b { 0 } else { 2 } {
66                            *c <<= 1;
67                        }
68                    }
69                }
70            }
71        }
72        Self { bit_width, mask }
73    }
74
75    pub fn builder<G>() -> SubmaskRangeQueryBuilder<G>
76    where
77        G: Group,
78    {
79        SubmaskRangeQueryBuilder::new()
80    }
81
82    pub fn get_query(&self, m: u32) -> impl Iterator<Item = (u32, bool)> {
83        let fix = m & self.mask[0];
84        let sub = m & self.mask[1];
85        let sup = (!m) & self.mask[2];
86        sup.subsets().flat_map(move |s| {
87            let inv = s.count_ones() & 1 == 1;
88            sub.subsets().map(move |t| (fix | s | t, inv))
89        })
90    }
91
92    pub fn update_query(&self, m: u32) -> impl Iterator<Item = u32> {
93        let fix = m & self.mask[0] | m & self.mask[1];
94        let sup = (!m) & self.mask[0];
95        let sub = m & self.mask[2];
96        sub.subsets()
97            .flat_map(move |s| sup.subsets().map(move |t| fix | s | t))
98    }
99}
100
101#[derive(Debug, Clone)]
102enum Query<T> {
103    Get { m: u32 },
104    Update { m: u32, x: T },
105}
106
107pub struct SubmaskRangeQueryBuilder<G>
108where
109    G: Group,
110{
111    query: Vec<Query<G::T>>,
112}
113
114impl<G> Debug for SubmaskRangeQueryBuilder<G>
115where
116    G: Group<T: Debug>,
117{
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        f.debug_struct("SubmaskRangeQueryBuilder")
120            .field("query", &self.query)
121            .finish()
122    }
123}
124
125impl<G> Default for SubmaskRangeQueryBuilder<G>
126where
127    G: Group,
128{
129    fn default() -> Self {
130        Self {
131            query: Default::default(),
132        }
133    }
134}
135
136impl<G> SubmaskRangeQueryBuilder<G>
137where
138    G: Group,
139{
140    pub fn new() -> Self {
141        Default::default()
142    }
143
144    pub fn push_get(&mut self, m: u32) {
145        self.query.push(Query::Get { m });
146    }
147
148    pub fn push_update(&mut self, m: u32, x: G::T) {
149        self.query.push(Query::Update { m, x });
150    }
151
152    pub fn solve(self) -> Vec<G::T> {
153        let s = SubmaskRangeQuery::new_with_queries(self.query.iter().map(|q| match q {
154            Query::Get { m } => (QueryKind::Get, *m),
155            Query::Update { m, .. } => (QueryKind::Update, *m),
156        }));
157        let out_size = self
158            .query
159            .iter()
160            .filter(|q| matches!(q, Query::Get { .. }))
161            .count();
162        let mut out = Vec::with_capacity(out_size);
163        let mut data = vec![G::unit(); 1 << s.bit_width];
164        for q in self.query {
165            match q {
166                Query::Get { m } => {
167                    let mut f = G::unit();
168                    let mut g = G::unit();
169                    for (k, inv) in s.get_query(m) {
170                        if inv {
171                            G::operate_assign(&mut g, &data[k as usize]);
172                        } else {
173                            G::operate_assign(&mut f, &data[k as usize]);
174                        }
175                    }
176                    out.push(G::rinv_operate(&f, &g));
177                }
178                Query::Update { m, x } => {
179                    for k in s.update_query(m) {
180                        G::operate_assign(&mut data[k as usize], &x);
181                    }
182                }
183            }
184        }
185        out
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::algebra::AdditiveOperation;
193
194    #[test]
195    fn test_submask_range_query() {
196        const W: u32 = 16;
197        let mut rng = Xorshift::default();
198        let mut q = SubmaskRangeQuery::builder::<AdditiveOperation<i32>>();
199        let mut a = vec![0; 1 << W];
200        let mut exp = vec![];
201        for _ in 0..2000 {
202            if rng.gen_bool(0.5) {
203                let i = rng.rand((1 << W) as _) as u32;
204                let x = rng.rand(100) as i32;
205                q.push_update(i, x);
206                a[i as usize] += x;
207            } else {
208                let i = rng.rand((1 << W) as _) as u32;
209                q.push_get(i);
210                let mut x = 0;
211                for j in 0..1 << W {
212                    if (i & j) == j {
213                        x += a[j as usize];
214                    }
215                }
216                exp.push(x);
217            }
218        }
219        let ans = q.solve();
220        assert_eq!(ans, exp);
221    }
222
223    #[test]
224    fn test_submask_range_query_online() {
225        const W: u32 = 16;
226        let mut rng = Xorshift::default();
227        let q = SubmaskRangeQuery::new(W);
228        let mut a = vec![0; 1 << W];
229        let mut b = vec![0; 1 << W];
230        for _ in 0..2000 {
231            if rng.gen_bool(0.5) {
232                let i = rng.rand((1 << W) as _) as u32;
233                let x = rng.rand(100) as i32;
234                a[i as usize] += x;
235                for j in q.update_query(i) {
236                    b[j as usize] += x;
237                }
238            } else {
239                let i = rng.rand((1 << W) as _) as u32;
240                let mut x = 0;
241                for j in 0..1 << W {
242                    if (i & j) == j {
243                        x += a[j as usize];
244                    }
245                }
246                let mut y = 0;
247                for (j, inv) in q.get_query(i) {
248                    if inv {
249                        y -= b[j as usize];
250                    } else {
251                        y += b[j as usize];
252                    }
253                }
254                assert_eq!(x, y);
255            }
256        }
257    }
258}
competitive/data_structure/submask_range_query.rs

competitive/data_structure/
submask_range_query.rs
