competitive/graph/
grid.rs

1use super::{
2    Adjacencies, AdjacenciesWithValue, AdjacencyView, AdjacencyViewIterFromValue, GraphBase,
3    VIndexWithValue, VertexMap, VertexView, Vertices,
4};
5use std::{iter::Map, marker::PhantomData, ops::Range};
6
7#[derive(Debug, Clone, Copy)]
8pub struct GridGraph<A> {
9    pub height: usize,
10    pub width: usize,
11    _marker: PhantomData<fn() -> A>,
12}
13
14impl GridGraph<Adj4> {
15    pub fn new_adj4(height: usize, width: usize) -> Self {
16        Self::new(height, width)
17    }
18    pub fn adj4(&self, vid: (usize, usize)) -> GridAdjacency<'_, Adj4> {
19        GridAdjacency {
20            g: self,
21            xy: vid,
22            diter: GridDirectionIter::default(),
23            _marker: PhantomData,
24        }
25    }
26}
27impl GridGraph<Adj8> {
28    pub fn new_adj8(height: usize, width: usize) -> Self {
29        Self::new(height, width)
30    }
31    pub fn adj8(&self, vid: (usize, usize)) -> GridAdjacency<'_, Adj8> {
32        GridAdjacency {
33            g: self,
34            xy: vid,
35            diter: GridDirectionIter::default(),
36            _marker: PhantomData,
37        }
38    }
39}
40
41impl<A> GridGraph<A> {
42    pub fn new(height: usize, width: usize) -> Self {
43        Self {
44            height,
45            width,
46            _marker: PhantomData,
47        }
48    }
49    pub fn move_by_diff(&self, xy: (usize, usize), dxdy: (isize, isize)) -> Option<(usize, usize)> {
50        let nx = xy.0.wrapping_add(dxdy.0 as usize);
51        let ny = xy.1.wrapping_add(dxdy.1 as usize);
52        if nx < self.height && ny < self.width {
53            Some((nx, ny))
54        } else {
55            None
56        }
57    }
58    pub fn flat(&self, xy: (usize, usize)) -> usize {
59        xy.0 * self.width + xy.1
60    }
61    pub fn unflat(&self, pos: usize) -> (usize, usize) {
62        (pos / self.width, pos % self.width)
63    }
64}
65
66impl<A> GraphBase for GridGraph<A> {
67    type VIndex = (usize, usize);
68}
69
70impl<A> Vertices for GridGraph<A> {
71    type VIter<'g>
72        = GridVertices
73    where
74        A: 'g;
75    fn vertices(&self) -> Self::VIter<'_> {
76        GridVertices {
77            xrange: 0..self.height,
78            yrange: 0..self.width,
79        }
80    }
81}
82
83#[derive(Debug, Clone)]
84pub struct GridVertices {
85    xrange: Range<usize>,
86    yrange: Range<usize>,
87}
88
89impl Iterator for GridVertices {
90    type Item = (usize, usize);
91    fn next(&mut self) -> Option<Self::Item> {
92        if self.xrange.start >= self.xrange.end {
93            None
94        } else if let Some(ny) = self.yrange.next() {
95            Some((self.xrange.start, ny))
96        } else {
97            self.yrange.start = 0;
98            self.xrange.start += 1;
99            self.next()
100        }
101    }
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
105pub enum GridDirection {
106    U = 0isize,
107    L = 1isize,
108    R = 2isize,
109    D = 3isize,
110    UL = 4isize,
111    UR = 5isize,
112    DL = 6isize,
113    DR = 7isize,
114}
115
116impl GridDirection {
117    pub fn dxdy(self) -> (isize, isize) {
118        match self {
119            GridDirection::U => (-1, 0),
120            GridDirection::L => (0, -1),
121            GridDirection::R => (0, 1),
122            GridDirection::D => (1, 0),
123            GridDirection::UL => (-1, -1),
124            GridDirection::UR => (-1, 1),
125            GridDirection::DL => (1, -1),
126            GridDirection::DR => (1, 1),
127        }
128    }
129    pub fn ndxdy(self, d: usize) -> (isize, isize) {
130        let d = d as isize;
131        match self {
132            GridDirection::U => (-d, 0),
133            GridDirection::L => (0, -d),
134            GridDirection::R => (0, d),
135            GridDirection::D => (d, 0),
136            GridDirection::UL => (-d, -d),
137            GridDirection::UR => (-d, d),
138            GridDirection::DL => (d, -d),
139            GridDirection::DR => (d, d),
140        }
141    }
142}
143
144impl Adjacencies for GridGraph<Adj4> {
145    type AIndex = VIndexWithValue<(usize, usize), GridDirection>;
146    type AIter<'g> = Map<
147        GridAdjacency<'g, Adj4>,
148        fn(((usize, usize), GridDirection)) -> VIndexWithValue<(usize, usize), GridDirection>,
149    >;
150    fn adjacencies(&self, vid: Self::VIndex) -> Self::AIter<'_> {
151        self.adj4(vid).map(Into::into)
152    }
153}
154impl Adjacencies for GridGraph<Adj8> {
155    type AIndex = VIndexWithValue<(usize, usize), GridDirection>;
156    type AIter<'g> = Map<
157        GridAdjacency<'g, Adj8>,
158        fn(((usize, usize), GridDirection)) -> VIndexWithValue<(usize, usize), GridDirection>,
159    >;
160    fn adjacencies(&self, vid: Self::VIndex) -> Self::AIter<'_> {
161        self.adj8(vid).map(Into::into)
162    }
163}
164
165impl AdjacenciesWithValue<GridDirection> for GridGraph<Adj4> {
166    type AIndex = VIndexWithValue<(usize, usize), GridDirection>;
167    type AIter<'g> = Map<
168        GridAdjacency<'g, Adj4>,
169        fn(((usize, usize), GridDirection)) -> VIndexWithValue<(usize, usize), GridDirection>,
170    >;
171    fn adjacencies_with_value(&self, vid: Self::VIndex) -> Self::AIter<'_> {
172        self.adjacencies(vid)
173    }
174}
175impl AdjacenciesWithValue<GridDirection> for GridGraph<Adj8> {
176    type AIndex = VIndexWithValue<(usize, usize), GridDirection>;
177    type AIter<'g> = Map<
178        GridAdjacency<'g, Adj8>,
179        fn(((usize, usize), GridDirection)) -> VIndexWithValue<(usize, usize), GridDirection>,
180    >;
181    fn adjacencies_with_value(&self, vid: Self::VIndex) -> Self::AIter<'_> {
182        self.adjacencies(vid)
183    }
184}
185
186impl<'a, M, T> AdjacencyView<'a, M, T> for GridGraph<Adj4>
187where
188    M: 'a + Fn(GridDirection) -> T,
189{
190    type AViewIter<'g> = AdjacencyViewIterFromValue<'g, 'a, Self, M, GridDirection, T>;
191    fn aviews<'g>(&'g self, map: &'a M, vid: Self::VIndex) -> Self::AViewIter<'g> {
192        AdjacencyViewIterFromValue::new(self.adjacencies(vid), map)
193    }
194}
195impl<'a, M, T> AdjacencyView<'a, M, T> for GridGraph<Adj8>
196where
197    M: 'a + Fn(GridDirection) -> T,
198{
199    type AViewIter<'g> = AdjacencyViewIterFromValue<'g, 'a, Self, M, GridDirection, T>;
200    fn aviews<'g>(&'g self, map: &'a M, vid: Self::VIndex) -> Self::AViewIter<'g> {
201        AdjacencyViewIterFromValue::new(self.adjacencies(vid), map)
202    }
203}
204
205#[derive(Debug, Clone, Copy)]
206pub enum Adj4 {}
207#[derive(Debug, Clone, Copy)]
208pub enum Adj8 {}
209
210#[derive(Debug, Clone)]
211pub struct GridDirectionIter<A> {
212    dir: Option<GridDirection>,
213    _marker: PhantomData<fn() -> A>,
214}
215impl<A> Default for GridDirectionIter<A> {
216    fn default() -> Self {
217        Self {
218            dir: Some(GridDirection::U),
219            _marker: PhantomData,
220        }
221    }
222}
223
224impl Iterator for GridDirectionIter<Adj4> {
225    type Item = GridDirection;
226    fn next(&mut self) -> Option<Self::Item> {
227        if let Some(dir) = &mut self.dir {
228            let cdir = Some(*dir);
229            self.dir = match dir {
230                GridDirection::U => Some(GridDirection::L),
231                GridDirection::L => Some(GridDirection::R),
232                GridDirection::R => Some(GridDirection::D),
233                _ => None,
234            };
235            cdir
236        } else {
237            None
238        }
239    }
240}
241impl Iterator for GridDirectionIter<Adj8> {
242    type Item = GridDirection;
243    fn next(&mut self) -> Option<Self::Item> {
244        if let Some(dir) = &mut self.dir {
245            let cdir = Some(*dir);
246            self.dir = match dir {
247                GridDirection::U => Some(GridDirection::L),
248                GridDirection::L => Some(GridDirection::R),
249                GridDirection::R => Some(GridDirection::D),
250                GridDirection::D => Some(GridDirection::UL),
251                GridDirection::UL => Some(GridDirection::UR),
252                GridDirection::UR => Some(GridDirection::DL),
253                GridDirection::DL => Some(GridDirection::DR),
254                GridDirection::DR => None,
255            };
256            cdir
257        } else {
258            None
259        }
260    }
261}
262
263#[derive(Debug, Clone)]
264pub struct GridAdjacency<'g, A> {
265    g: &'g GridGraph<A>,
266    xy: (usize, usize),
267    diter: GridDirectionIter<A>,
268    _marker: PhantomData<fn() -> A>,
269}
270
271impl<A> Iterator for GridAdjacency<'_, A>
272where
273    GridDirectionIter<A>: Iterator<Item = GridDirection>,
274{
275    type Item = ((usize, usize), GridDirection);
276    fn next(&mut self) -> Option<Self::Item> {
277        for dir in self.diter.by_ref() {
278            match self.g.move_by_diff(self.xy, dir.dxdy()) {
279                Some(nxy) => return Some((nxy, dir)),
280                None => continue,
281            }
282        }
283        None
284    }
285}
286
287impl<A, T> VertexMap<T> for GridGraph<A> {
288    type Vmap = Vec<Vec<T>>;
289    fn construct_vmap<F>(&self, mut f: F) -> Self::Vmap
290    where
291        F: FnMut() -> T,
292    {
293        (0..self.height)
294            .map(|_| (0..self.width).map(|_| f()).collect())
295            .collect()
296    }
297    fn vmap_get<'a>(&self, map: &'a Self::Vmap, (x, y): Self::VIndex) -> &'a T {
298        assert!(x < self.height, "expected 0..{}, but {}", self.height, x);
299        assert!(y < self.width, "expected 0..{}, but {}", self.width, y);
300        unsafe { map.get_unchecked(x).get_unchecked(y) }
301    }
302    fn vmap_get_mut<'a>(&self, map: &'a mut Self::Vmap, (x, y): Self::VIndex) -> &'a mut T {
303        assert!(x < self.height, "expected 0..{}, but {}", self.height, x);
304        assert!(y < self.width, "expected 0..{}, but {}", self.width, y);
305        unsafe { map.get_unchecked_mut(x).get_unchecked_mut(y) }
306    }
307}
308impl<A, T> VertexView<Vec<Vec<T>>, T> for GridGraph<A>
309where
310    T: Clone,
311{
312    fn vview(&self, map: &Vec<Vec<T>>, vid: Self::VIndex) -> T {
313        self.vmap_get(map, vid).clone()
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::GridGraph;
320    use crate::{
321        algebra::AdditiveOperation,
322        graph::{ShortestPathExt, StandardSp},
323        num::Saturating,
324        tools::Xorshift,
325    };
326
327    #[test]
328    fn grid_graph_apsp() {
329        let mut rng = Xorshift::default();
330        const A: u64 = 1_000_000_000;
331        let h = rng.rand(15) as usize + 1;
332        let w = rng.rand(15) as usize + 1;
333
334        let weight: Vec<_> = std::iter::repeat_with(|| Saturating(rng.rand(A - 1) + 1))
335            .take(8)
336            .collect();
337        type Sp = StandardSp<AdditiveOperation<Saturating<u64>>>;
338
339        let g = GridGraph::new_adj4(h, w);
340        let cost: Vec<Vec<Vec<Vec<_>>>> = (0..h)
341            .map(|i| {
342                (0..w)
343                    .map(|j| g.dijkstra_ss::<Sp, _>((i, j), &|dir| weight[dir as usize]))
344                    .collect()
345            })
346            .collect();
347        let cost2: Vec<Vec<_>> = g.warshall_floyd_ap::<Sp, _>(&|dir| weight[dir as usize]);
348        assert_eq!(cost, cost2);
349
350        let g = GridGraph::new_adj8(h, w);
351        let cost: Vec<Vec<Vec<Vec<_>>>> = (0..h)
352            .map(|i| {
353                (0..w)
354                    .map(|j| g.dijkstra_ss::<Sp, _>((i, j), &|dir| weight[dir as usize]))
355                    .collect()
356            })
357            .collect();
358        let cost2: Vec<Vec<_>> = g.warshall_floyd_ap::<Sp, _>(&|dir| weight[dir as usize]);
359        assert_eq!(cost, cost2);
360    }
361}