Skip to main content

nodedb_spatial/rtree/
split.rs

1//! R*-tree node splitting (axis split strategy).
2
3use super::node::{ChildRef, MIN_FILL_INTERNAL, MIN_FILL_LEAF, Node, NodeKind, RTreeEntry};
4use super::tree::RTree;
5
6/// Split an overflowing node using the R*-tree axis split strategy.
7pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) {
8    let is_root = node_idx == tree.root;
9    let level = tree.nodes[node_idx].level;
10
11    let (sibling_idx, sibling_bbox) = match &mut tree.nodes[node_idx].kind {
12        NodeKind::Leaf { entries } => {
13            let all = std::mem::take(entries);
14            let (keep, split_off) = split_leaf_entries(all);
15            if let NodeKind::Leaf { entries } = &mut tree.nodes[node_idx].kind {
16                *entries = keep;
17            }
18            tree.nodes[node_idx].recompute_bbox();
19
20            let mut sibling = Node::new_leaf();
21            if let NodeKind::Leaf { entries } = &mut sibling.kind {
22                *entries = split_off;
23            }
24            sibling.recompute_bbox();
25            let bbox = sibling.bbox;
26            let idx = tree.nodes.len();
27            tree.nodes.push(sibling);
28            (idx, bbox)
29        }
30        NodeKind::Internal { children } => {
31            let all = std::mem::take(children);
32            let (keep, split_off) = split_internal_children(all);
33            if let NodeKind::Internal { children } = &mut tree.nodes[node_idx].kind {
34                *children = keep;
35            }
36            tree.nodes[node_idx].recompute_bbox();
37
38            let mut sibling = Node::new_internal(level);
39            if let NodeKind::Internal { children } = &mut sibling.kind {
40                *children = split_off;
41            }
42            sibling.recompute_bbox();
43            let bbox = sibling.bbox;
44            let idx = tree.nodes.len();
45            tree.nodes.push(sibling);
46            (idx, bbox)
47        }
48    };
49
50    if is_root {
51        let old_root_bbox = tree.nodes[node_idx].bbox;
52        let mut new_root = Node::new_internal(level + 1);
53        if let NodeKind::Internal { children } = &mut new_root.kind {
54            children.push(ChildRef {
55                bbox: old_root_bbox,
56                node_idx,
57            });
58            children.push(ChildRef {
59                bbox: sibling_bbox,
60                node_idx: sibling_idx,
61            });
62        }
63        new_root.recompute_bbox();
64        let new_root_idx = tree.nodes.len();
65        tree.nodes.push(new_root);
66        tree.root = new_root_idx;
67    } else {
68        let parent_idx = tree.find_parent(tree.root, node_idx);
69        if let Some(pidx) = parent_idx {
70            let updated_bbox = tree.nodes[node_idx].bbox;
71            if let NodeKind::Internal { children } = &mut tree.nodes[pidx].kind {
72                children.push(ChildRef {
73                    bbox: sibling_bbox,
74                    node_idx: sibling_idx,
75                });
76                for c in children.iter_mut() {
77                    if c.node_idx == node_idx {
78                        c.bbox = updated_bbox;
79                        break;
80                    }
81                }
82            }
83            tree.nodes[pidx].recompute_bbox();
84            if tree.nodes[pidx].is_overflow() {
85                split_node(tree, pidx);
86            }
87        }
88    }
89}
90
91/// Choose split axis (min margin sum), then split index (min overlap).
92fn split_leaf_entries(mut entries: Vec<RTreeEntry>) -> (Vec<RTreeEntry>, Vec<RTreeEntry>) {
93    let min_fill = MIN_FILL_LEAF;
94    let best_axis = choose_best_axis_leaf(&mut entries, min_fill);
95
96    sort_entries_by_axis(&mut entries, best_axis);
97    let best_k = choose_best_split_leaf(&entries, min_fill);
98
99    let split_off = entries.split_off(best_k);
100    (entries, split_off)
101}
102
103fn split_internal_children(mut children: Vec<ChildRef>) -> (Vec<ChildRef>, Vec<ChildRef>) {
104    let min_fill = MIN_FILL_INTERNAL;
105    let best_axis = choose_best_axis_internal(&mut children, min_fill);
106
107    sort_children_by_axis(&mut children, best_axis);
108    let best_k = choose_best_split_internal(&children, min_fill);
109
110    let split_off = children.split_off(best_k);
111    (children, split_off)
112}
113
114fn choose_best_axis_leaf(entries: &mut [RTreeEntry], min_fill: usize) -> usize {
115    let mut best_axis = 0;
116    let mut best_margin = f64::INFINITY;
117    for axis in 0..2 {
118        sort_entries_by_axis(entries, axis);
119        let mut margin_sum = 0.0;
120        for k in min_fill..=(entries.len() - min_fill) {
121            let left = entries[..k]
122                .iter()
123                .fold(entries[0].bbox, |a, e| a.union(&e.bbox));
124            let right = entries[k..]
125                .iter()
126                .fold(entries[k].bbox, |a, e| a.union(&e.bbox));
127            margin_sum += left.margin() + right.margin();
128        }
129        if margin_sum < best_margin {
130            best_margin = margin_sum;
131            best_axis = axis;
132        }
133    }
134    best_axis
135}
136
137fn choose_best_axis_internal(children: &mut [ChildRef], min_fill: usize) -> usize {
138    let mut best_axis = 0;
139    let mut best_margin = f64::INFINITY;
140    for axis in 0..2 {
141        sort_children_by_axis(children, axis);
142        let mut margin_sum = 0.0;
143        for k in min_fill..=(children.len() - min_fill) {
144            let left = children[..k]
145                .iter()
146                .fold(children[0].bbox, |a, c| a.union(&c.bbox));
147            let right = children[k..]
148                .iter()
149                .fold(children[k].bbox, |a, c| a.union(&c.bbox));
150            margin_sum += left.margin() + right.margin();
151        }
152        if margin_sum < best_margin {
153            best_margin = margin_sum;
154            best_axis = axis;
155        }
156    }
157    best_axis
158}
159
160fn choose_best_split_leaf(entries: &[RTreeEntry], min_fill: usize) -> usize {
161    let mut best_k = min_fill;
162    let mut best_overlap = f64::INFINITY;
163    let mut best_area = f64::INFINITY;
164    for k in min_fill..=(entries.len() - min_fill) {
165        let left = entries[..k]
166            .iter()
167            .fold(entries[0].bbox, |a, e| a.union(&e.bbox));
168        let right = entries[k..]
169            .iter()
170            .fold(entries[k].bbox, |a, e| a.union(&e.bbox));
171        let overlap = left.overlap_area(&right);
172        let area = left.area() + right.area();
173        if overlap < best_overlap || (overlap == best_overlap && area < best_area) {
174            best_overlap = overlap;
175            best_area = area;
176            best_k = k;
177        }
178    }
179    best_k
180}
181
182fn choose_best_split_internal(children: &[ChildRef], min_fill: usize) -> usize {
183    let mut best_k = min_fill;
184    let mut best_overlap = f64::INFINITY;
185    let mut best_area = f64::INFINITY;
186    for k in min_fill..=(children.len() - min_fill) {
187        let left = children[..k]
188            .iter()
189            .fold(children[0].bbox, |a, c| a.union(&c.bbox));
190        let right = children[k..]
191            .iter()
192            .fold(children[k].bbox, |a, c| a.union(&c.bbox));
193        let overlap = left.overlap_area(&right);
194        let area = left.area() + right.area();
195        if overlap < best_overlap || (overlap == best_overlap && area < best_area) {
196            best_overlap = overlap;
197            best_area = area;
198            best_k = k;
199        }
200    }
201    best_k
202}
203
204fn sort_entries_by_axis(entries: &mut [RTreeEntry], axis: usize) {
205    entries.sort_by(|a, b| {
206        let va = if axis == 0 {
207            a.bbox.min_lng
208        } else {
209            a.bbox.min_lat
210        };
211        let vb = if axis == 0 {
212            b.bbox.min_lng
213        } else {
214            b.bbox.min_lat
215        };
216        va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal)
217    });
218}
219
220fn sort_children_by_axis(children: &mut [ChildRef], axis: usize) {
221    children.sort_by(|a, b| {
222        let va = if axis == 0 {
223            a.bbox.min_lng
224        } else {
225            a.bbox.min_lat
226        };
227        let vb = if axis == 0 {
228            b.bbox.min_lng
229        } else {
230            b.bbox.min_lat
231        };
232        va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal)
233    });
234}