101 lines
3 KiB
Rust
101 lines
3 KiB
Rust
use itertools::Itertools;
|
|
|
|
#[derive(Debug)]
|
|
pub struct HeightMap {
|
|
pub map: Vec<Vec<usize>>,
|
|
pub basin: Vec<Vec<Option<usize>>>,
|
|
}
|
|
|
|
impl HeightMap {
|
|
pub fn new(map: Vec<Vec<usize>>) -> Self {
|
|
let mut basin = Vec::new();
|
|
for line in &map {
|
|
let mut basin_line = Vec::new();
|
|
for _point in line {
|
|
basin_line.push(None)
|
|
}
|
|
basin.push(basin_line);
|
|
}
|
|
let mut new_self = Self { map, basin };
|
|
let mut basin_counter = 0;
|
|
for y in 0..new_self.map.len() {
|
|
for x in 0..new_self.map[y].len() {
|
|
if new_self.map[y][x] < 9 && new_self.basin[y][x].is_none() {
|
|
new_self.flood_basin((x, y), basin_counter);
|
|
basin_counter += 1;
|
|
}
|
|
}
|
|
}
|
|
new_self
|
|
}
|
|
|
|
pub fn get_low_points(&self) -> Vec<(isize, isize)> {
|
|
let mut low_points = Vec::new();
|
|
for y in 0..self.map.len() {
|
|
for x in 0..self.map[y].len() {
|
|
if self.is_low_point((x as isize, y as isize)) {
|
|
low_points.push((x as isize, y as isize))
|
|
}
|
|
}
|
|
}
|
|
low_points
|
|
}
|
|
|
|
pub fn is_low_point(&self, position: (isize, isize)) -> bool {
|
|
let val = match self.get(position) {
|
|
Some(pos) => pos,
|
|
None => return false,
|
|
};
|
|
let (x, y) = position;
|
|
for (x_off, y_off) in [(1isize, 0isize), (0, 1), (-1, 0), (0, -1)].iter() {
|
|
if let Some(neighbour) = self.get((x + x_off, y + y_off)) {
|
|
if neighbour <= val {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
true
|
|
}
|
|
|
|
pub fn flood_basin(&mut self, (x, y): (usize, usize), number: usize) {
|
|
if let Some(basin) = self.basin[y][x] {
|
|
assert_eq!(basin, number);
|
|
return;
|
|
} else {
|
|
self.basin[y][x] = Some(number)
|
|
}
|
|
for (x_off, y_off) in [(1isize, 0isize), (0, 1), (-1, 0), (0, -1)].iter() {
|
|
if let Some(neighbour) = self.get((x as isize + x_off, y as isize + y_off)) {
|
|
if neighbour < 9 {
|
|
self.flood_basin(
|
|
((x as isize + x_off) as usize, (y as isize + y_off) as usize),
|
|
number,
|
|
)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn get(&self, (x, y): (isize, isize)) -> Option<usize> {
|
|
let (x, y): (usize, usize) = match (x.try_into(), y.try_into()) {
|
|
(Ok(x), Ok(y)) => (x, y),
|
|
_ => return None,
|
|
};
|
|
self.map.get(y).map(|val| val.get(x)).flatten().copied()
|
|
}
|
|
|
|
pub fn get_risk_level(&self, position: (isize, isize)) -> Option<usize> {
|
|
self.get(position).map(|x| x + 1)
|
|
}
|
|
|
|
pub fn get_basin_sizes(&self) -> Vec<usize> {
|
|
self.basin
|
|
.iter()
|
|
.flat_map(|line| line.iter())
|
|
.filter_map(|val| *val)
|
|
.counts()
|
|
.into_values()
|
|
.collect()
|
|
}
|
|
}
|