h3ron/algorithm/
cell_clusters.rs

1use crate::collections::HashMap;
2use crate::iter::GridDiskBuilder;
3use crate::{Error, H3Cell};
4use ahash::RandomState;
5use hashbrown::hash_map::Entry;
6use indexmap::IndexMap;
7use std::cmp::Ordering;
8
9/// find clusters of neighboring cells
10///
11/// Requires the `indexmap` feature.
12pub fn find_cell_clusters<CellsIter>(cells: CellsIter) -> Result<Vec<Vec<H3Cell>>, Error>
13where
14    CellsIter: Iterator<Item = H3Cell>,
15{
16    Ok(find_cell_clusters_eq_value_impl(cells)?
17        .into_values()
18        .map(|(cluster, _)| cluster)
19        .collect())
20}
21
22/// find clusters of neighboring cells where the same value is associated with the cells.
23///
24/// Cells are assumed to be unique, otherwise the behaviour is undefined.
25///
26/// Requires the `indexmap` feature.
27pub fn find_cell_clusters_eq_value<CellValueIter, CV, Value>(
28    cell_value_iter: CellValueIter,
29) -> Result<Vec<(Vec<H3Cell>, Value)>, Error>
30where
31    CV: CellAndValue<Value>,
32    CellValueIter: Iterator<Item = CV>,
33    Value: PartialEq,
34{
35    Ok(find_cell_clusters_eq_value_impl(cell_value_iter)?
36        .into_values()
37        .collect())
38}
39
40pub trait CellAndValue<Value> {
41    fn cell(&self) -> H3Cell;
42    fn value(self) -> Value;
43}
44
45impl CellAndValue<()> for H3Cell {
46    fn cell(&self) -> H3Cell {
47        *self
48    }
49
50    fn value(self) {}
51}
52
53impl<Value> CellAndValue<Value> for (H3Cell, Value) {
54    fn cell(&self) -> H3Cell {
55        self.0
56    }
57
58    fn value(self) -> Value {
59        self.1
60    }
61}
62
63fn find_cell_clusters_eq_value_impl<CellValueIter, CV, Value>(
64    cell_value_iter: CellValueIter,
65) -> Result<HashMap<usize, (Vec<H3Cell>, Value)>, Error>
66where
67    CV: CellAndValue<Value>,
68    CellValueIter: Iterator<Item = CV>,
69    Value: PartialEq,
70{
71    let items: IndexMap<_, _, RandomState> =
72        cell_value_iter.map(|cv| (cv.cell(), cv.value())).collect();
73    let mut cluster_ids: Vec<usize> = (0..items.len()).collect();
74
75    let mut mutated = true;
76    let mut disk_builder = GridDiskBuilder::create(1, 1)?;
77    while mutated {
78        mutated = false;
79        for (pos, (cell, value)) in items.iter().enumerate() {
80            let mut least_cluster_id = cluster_ids[pos];
81            for (neighbor_cell, _) in &mut disk_builder.build_grid_disk(cell)? {
82                if let Some((neighbor_pos, _, neighbor_value)) = items.get_full(&neighbor_cell) {
83                    if neighbor_value == value {
84                        match cluster_ids[neighbor_pos].cmp(&least_cluster_id) {
85                            Ordering::Less => {
86                                least_cluster_id = cluster_ids[neighbor_pos];
87                                cluster_ids[pos] = least_cluster_id;
88                                mutated = true;
89                            }
90                            Ordering::Equal => {}
91                            Ordering::Greater => {
92                                cluster_ids[neighbor_pos] = least_cluster_id;
93                                mutated = true;
94                            }
95                        }
96                    }
97                }
98            }
99        }
100    }
101
102    Ok(cluster_ids.into_iter().zip(items).fold(
103        HashMap::default(),
104        |mut acc, (group, (cell, value))| {
105            match acc.entry(group) {
106                Entry::Vacant(entry) => {
107                    entry.insert((vec![cell], value));
108                }
109                Entry::Occupied(mut entry) => {
110                    entry.get_mut().0.push(cell);
111                }
112            };
113            acc
114        },
115    ))
116}
117
118#[cfg(test)]
119mod tests {
120    use crate::algorithm::find_cell_clusters;
121    use crate::H3Cell;
122
123    #[test]
124    fn find_cell_clusters_simple() {
125        let mut disk1: Vec<_> = H3Cell::from_coordinate((12.2, 14.5).into(), 6)
126            .unwrap()
127            .grid_disk(3)
128            .unwrap()
129            .iter()
130            .collect();
131        disk1.sort_unstable();
132        let mut disk2: Vec<_> = H3Cell::from_coordinate((42.2, 45.5).into(), 6)
133            .unwrap()
134            .grid_disk(2)
135            .unwrap()
136            .iter()
137            .collect();
138        disk2.sort_unstable();
139
140        let mut clusters =
141            find_cell_clusters(disk1.iter().copied().chain(disk2.iter().copied())).unwrap();
142        assert_eq!(clusters.len(), 2);
143        let mut cluster1 = clusters.remove(0);
144        cluster1.sort_unstable();
145        let mut cluster2 = clusters.remove(0);
146        cluster2.sort_unstable();
147        assert!(cluster1 == disk1 || cluster1 == disk2);
148        assert!(cluster2 == disk1 || cluster2 == disk2);
149        assert_ne!(cluster1, cluster2);
150    }
151}