impl<T: Ord + From<u8> + PartialEq<T>> Grid<T> {
fn basins(&self) -> Vec<((usize, usize), usize)> {
let mut uf = self
.map_with_coords(|p, n| (p, if n == &T::from(9) { 0 } else { 1 }));
for (p, h) in uf
.points()
.zip(self.inner.iter())
.filter(|(_, h)| *h < &T::from(9))
{
// I originally thought this would also find the minimum of each
// basin but it doesn't always.
for (n, j) in uf
.neighbours(p)
.map(|n| (n, &self[n]))
.filter(|(_, n)| *n < &T::from(9))
{
match h.cmp(j) {
std::cmp::Ordering::Less => uf.flow(n, p),
std::cmp::Ordering::Equal => (),
std::cmp::Ordering::Greater => uf.flow(p, n),
}
}
}
uf.points()
.zip(uf.inner.iter().map(|(_, s)| s))
.filter(|(_, s)| *s > &0)
.map(|(p, s)| (p, *s))
.collect()
}
}
impl<T: Copy> Grid<((usize, usize), T)> {
fn u_find(&mut self, p: (usize, usize)) -> (usize, usize) {
let d = unsafe { &mut *(&mut self[p].0 as *mut (usize, usize)) };
if d == &p {
p
} else {
let q = self.u_find(*d);
*d = q;
q
}
}
}
impl<T: Copy + From<u8> + std::ops::Add<T, Output = T>>
Grid<((usize, usize), T)>
{
fn flow(&mut self, mut s: (usize, usize), mut t: (usize, usize)) {
s = self.u_find(s);
t = self.u_find(t);
if s != t {
let y = unsafe { &mut *(&mut self[t].1 as *mut T) };
let (d, x) = &mut self[s];
let x = std::mem::replace(x, T::from(0));
*y = *y + x;
*d = t;
}
}
}