static_bushes/kdbush/
sort.rs

1use crate::{kdbush::AllowedNumber, util::IndexVec};
2
3pub fn sort_kd<T: AllowedNumber>(
4    ids: &mut IndexVec,
5    coords: &mut [T],
6    node_size: usize,
7    left: usize,
8    right: usize,
9    axis: usize,
10) {
11    if right - left <= node_size {
12        return;
13    }
14
15    let m = (left + right) >> 1; // middle index
16
17    // sort ids and coords around the middle index so that the halves lie
18    // either left/right or top/bottom correspondingly (taking turns)
19    select(ids, coords, m, left, right, axis);
20
21    // recursively kd-sort first half and second half on the opposite axis
22    sort_kd(ids, coords, node_size, left, m - 1, 1 - axis);
23    sort_kd(ids, coords, node_size, m + 1, right, 1 - axis);
24}
25
26// custom Floyd-Rivest selection algorithm: sort ids and coords so that
27// [left..k-1] items are smaller than k-th item (on either x or y axis)
28fn select<T: AllowedNumber>(
29    ids: &mut IndexVec,
30    coords: &mut [T],
31    k: usize,
32    mut left: usize,
33    mut right: usize,
34    axis: usize,
35) {
36    while right > left {
37        if right - left > 600 {
38            let n = (right - left + 1) as f64;
39            let m = (k - left + 1) as f64;
40            let fk = k as f64;
41
42            let z = n.ln();
43            let s = 0.5 * (2.0 * z / 3.0).exp();
44            let sd =
45                0.5 * (z * s * (n - s) / n).sqrt() * (if m - n / 2.0 < 0.0 { -1.0 } else { 1.0 });
46            let new_left = max(left, (fk - m * s / n + sd).floor() as usize);
47            let new_right = min(right, (fk + (n - m) * s / n + sd).floor() as usize);
48            select(ids, coords, k, new_left, new_right, axis);
49        }
50
51        let t = coords[2 * k + axis];
52        let mut i = left;
53        let mut j = right;
54
55        swap_item(ids, coords, left, k);
56        if coords[2 * right + axis] > t {
57            swap_item(ids, coords, left, right);
58        }
59
60        while i < j {
61            swap_item(ids, coords, i, j);
62            i += 1;
63            j -= 1;
64            while coords[2 * i + axis] < t {
65                i += 1
66            }
67            while coords[2 * j + axis] > t {
68                j -= 1
69            }
70        }
71
72        if coords[2 * left + axis] == t {
73            swap_item(ids, coords, left, j)
74        } else {
75            j += 1;
76            swap_item(ids, coords, j, right);
77        }
78
79        if j <= k {
80            left = j + 1;
81        }
82        if k <= j {
83            right = j - 1;
84        }
85    }
86}
87
88fn swap_item<T: AllowedNumber>(ids: &mut IndexVec, coords: &mut [T], i: usize, j: usize) {
89    ids.swap(i, j);
90    coords.swap(2 * i, 2 * j);
91    coords.swap(2 * i + 1, 2 * j + 1);
92}
93
94// implementing these myself to make the library work with floats even though they're not
95// Ord; eventually it would be better to make OrderedFloat and num_traits play nice
96#[inline(always)]
97fn min<U: PartialOrd>(a: U, b: U) -> U {
98    if a <= b {
99        a
100    } else {
101        b
102    }
103}
104
105#[inline(always)]
106fn max<U: PartialOrd>(a: U, b: U) -> U {
107    if b >= a {
108        b
109    } else {
110        a
111    }
112}
static_bushes/kdbush/sort.rs

static_bushes/kdbush/
sort.rs
