Skip to main content

nodedb_spatial/rtree/
split.rs

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