h3ron/algorithm/
cell_clusters.rs1use 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
9pub 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
22pub 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}