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}