scirs2_spatial/rtree/

query.rs

use crate::error::SpatialResult;
use crate::rtree::node::{Entry, EntryWithDistance, Node, RTree, Rectangle};
use scirs2_core::ndarray::ArrayView1;
use std::cmp::Ordering;
use std::collections::BinaryHeap;

impl<T: Clone> RTree<T> {
    /// Search for data points within a range
    ///
    /// # Arguments
    ///
    /// * `min` - Minimum coordinates of the search range
    /// * `max` - Maximum coordinates of the search range
    ///
    /// # Returns
    ///
    /// A `SpatialResult` containing a vector of (index, data) pairs for data points within the range,
    /// or an error if the range has invalid dimensions
    pub fn search_range(
        &self,
        min: &ArrayView1<f64>,
        max: &ArrayView1<f64>,
    ) -> SpatialResult<Vec<(usize, T)>> {
        if min.len() != self.ndim() || max.len() != self.ndim() {
            return Err(crate::error::SpatialError::DimensionError(format!(
                "Search range dimensions ({}, {}) do not match RTree dimension {}",
                min.len(),
                max.len(),
                self.ndim()
            )));
        }

        // Create a search rectangle
        let rect = Rectangle::new(min.to_owned(), max.to_owned())?;

        // Perform the search
        let mut results = Vec::new();
        self.search_range_internal(&rect, &self.root, &mut results)?;

        Ok(results)
    }

    /// Recursively search for points within a range
    #[allow(clippy::only_used_in_recursion)]
    fn search_range_internal(
        &self,
        rect: &Rectangle,
        node: &Node<T>,
        results: &mut Vec<(usize, T)>,
    ) -> SpatialResult<()> {
        // Process each entry in the node
        for entry in &node.entries {
            // Check if this entry's MBR intersects with the search rectangle
            if entry.mbr().intersects(rect)? {
                match entry {
                    // If this is a leaf entry, add the data to the results
                    Entry::Leaf { data, index, .. } => {
                        results.push((*index, data.clone()));
                    }
                    // If this is a non-leaf entry, recursively search its child
                    Entry::NonLeaf { child, .. } => {
                        self.search_range_internal(rect, child, results)?;
                    }
                }
            }
        }

        Ok(())
    }

    /// Find the k nearest neighbors to a query point
    ///
    /// # Arguments
    ///
    /// * `point` - The query point
    /// * `k` - The number of nearest neighbors to find
    ///
    /// # Returns
    ///
    /// A `SpatialResult` containing a vector of (index, data, distance) tuples for the k nearest data points,
    /// sorted by distance (closest first), or an error if the point has invalid dimensions
    pub fn nearest(
        &self,
        point: &ArrayView1<f64>,
        k: usize,
    ) -> SpatialResult<Vec<(usize, T, f64)>> {
        if point.len() != self.ndim() {
            return Err(crate::error::SpatialError::DimensionError(format!(
                "Point dimension {} does not match RTree dimension {}",
                point.len(),
                self.ndim()
            )));
        }

        if k == 0 || self.is_empty() {
            return Ok(Vec::new());
        }

        // Use a priority queue to keep track of nodes to visit
        let mut pq = BinaryHeap::new();
        let mut results = Vec::new();

        // Initialize with root node
        if let Ok(Some(root_mbr)) = self.root.mbr() {
            let _distance = root_mbr.min_distance_to_point(point)?;

            // Add all entries from the root
            for entry in &self.root.entries {
                let entry_distance = entry.mbr().min_distance_to_point(point)?;
                pq.push(EntryWithDistance {
                    entry: entry.clone(),
                    distance: entry_distance,
                });
            }
        }

        // Current maximum distance in the result set
        let mut max_distance = f64::MAX;

        // Process the priority queue
        while let Some(item) = pq.pop() {
            // If the minimum distance is greater than our current maximum, we can stop
            if item.distance > max_distance && results.len() >= k {
                break;
            }

            match item.entry {
                // If this is a leaf entry, add it to the results
                Entry::Leaf { data, index, .. } => {
                    results.push((index, data, item.distance));

                    // Update max_distance if we have enough results
                    if results.len() >= k {
                        // Sort results by distance
                        results.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(Ordering::Equal));

                        // Keep only the k closest
                        results.truncate(k);

                        // Update max_distance
                        if let Some((_, _, dist)) = results.last() {
                            max_distance = *dist;
                        }
                    }
                }
                // If this is a non-leaf entry, add its children to the queue
                Entry::NonLeaf { child, .. } => {
                    for entry in &child.entries {
                        let entry_distance = entry.mbr().min_distance_to_point(point)?;

                        // Only add entries that could be closer than our current maximum
                        if entry_distance <= max_distance || results.len() < k {
                            pq.push(EntryWithDistance {
                                entry: entry.clone(),
                                distance: entry_distance,
                            });
                        }
                    }
                }
            }
        }

        // Sort final results by distance
        results.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(Ordering::Equal));

        // Truncate to k results
        results.truncate(k);

        Ok(results)
    }

    /// Perform a spatial join between this R-tree and another
    ///
    /// # Arguments
    ///
    /// * `other` - The other R-tree to join with
    /// * `predicate` - A function that takes MBRs from both trees and returns true
    ///   if they should be joined, e.g., for an intersection join: `|mbr1, mbr2| mbr1.intersects(mbr2)`
    ///
    /// # Returns
    ///
    /// A `SpatialResult` containing a vector of pairs of data from both trees that satisfy the predicate,
    /// or an error if the R-trees have different dimensions
    pub fn spatial_join<U, P>(&self, other: &RTree<U>, predicate: P) -> SpatialResult<Vec<(T, U)>>
    where
        U: Clone,
        P: Fn(&Rectangle, &Rectangle) -> SpatialResult<bool>,
    {
        if self.ndim() != other.ndim() {
            return Err(crate::error::SpatialError::DimensionError(format!(
                "RTrees have different dimensions: {} and {}",
                self.ndim(),
                other.ndim()
            )));
        }

        let mut results = Vec::new();

        // If either tree is empty, return an empty result
        if self.is_empty() || other.is_empty() {
            return Ok(results);
        }

        // Perform the join
        self.spatial_join_internal(&self.root, &other.root, &predicate, &mut results)?;

        Ok(results)
    }

    /// Recursively perform a spatial join between two nodes
    #[allow(clippy::only_used_in_recursion)]
    fn spatial_join_internal<U, P>(
        &self,
        node1: &Node<T>,
        node2: &Node<U>,
        predicate: &P,
        results: &mut Vec<(T, U)>,
    ) -> SpatialResult<()>
    where
        U: Clone,
        P: Fn(&Rectangle, &Rectangle) -> SpatialResult<bool>,
    {
        // Process each pair of entries
        for entry1 in &node1.entries {
            for entry2 in &node2.entries {
                // Check if the entries satisfy the predicate
                if predicate(entry1.mbr(), entry2.mbr())? {
                    match (entry1, entry2) {
                        // If both are leaf entries, add to results
                        (Entry::Leaf { data: data1, .. }, Entry::Leaf { data: data2, .. }) => {
                            results.push((data1.clone(), data2.clone()));
                        }
                        // If entry1 is a non-leaf, recurse with its children
                        (Entry::NonLeaf { child: child1, .. }, Entry::Leaf { .. }) => {
                            self.spatial_join_internal(
                                child1,
                                &Node {
                                    entries: vec![entry2.clone()],
                                    _isleaf: true,
                                    level: 0,
                                },
                                predicate,
                                results,
                            )?;
                        }
                        // If entry2 is a non-leaf, recurse with its children
                        (Entry::Leaf { .. }, Entry::NonLeaf { child: child2, .. }) => {
                            self.spatial_join_internal(
                                &Node {
                                    entries: vec![entry1.clone()],
                                    _isleaf: true,
                                    level: 0,
                                },
                                child2,
                                predicate,
                                results,
                            )?;
                        }
                        // If both are non-leaf entries, recurse with both children
                        (
                            Entry::NonLeaf { child: child1, .. },
                            Entry::NonLeaf { child: child2, .. },
                        ) => {
                            self.spatial_join_internal(child1, child2, predicate, results)?;
                        }
                    }
                }
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_rtree_nearest_neighbors() {
        // Create a new R-tree
        let mut rtree: RTree<i32> = RTree::new(2, 2, 4).unwrap();

        // Insert some points
        let points = vec![
            (array![0.0, 0.0], 0),
            (array![1.0, 0.0], 1),
            (array![0.0, 1.0], 2),
            (array![1.0, 1.0], 3),
            (array![0.5, 0.5], 4),
            (array![2.0, 2.0], 5),
            (array![3.0, 3.0], 6),
            (array![4.0, 4.0], 7),
            (array![5.0, 5.0], 8),
            (array![6.0, 6.0], 9),
        ];

        for (point, value) in points {
            rtree.insert(point, value).unwrap();
        }

        // Find the nearest neighbor to (0.6, 0.6)
        let nn_results = rtree.nearest(&array![0.6, 0.6].view(), 1).unwrap();

        // Should be (0.5, 0.5)
        assert_eq!(nn_results.len(), 1);
        assert_eq!(nn_results[0].1, 4);

        // Find the 3 nearest neighbors to (0.0, 0.0)
        let nn_results = rtree.nearest(&array![0.0, 0.0].view(), 3).unwrap();

        // Should be (0.0, 0.0), (1.0, 0.0), and (0.0, 1.0)
        assert_eq!(nn_results.len(), 3);

        // The results should be sorted by distance
        assert_eq!(nn_results[0].1, 0); // (0.0, 0.0) - distance 0
        assert_eq!(nn_results[1].1, 4); // (0.5, 0.5) - distance ~0.707

        // The third one could be either (1.0, 0.0) or (0.0, 1.0) - distance 1.0
        assert!(nn_results[2].1 == 1 || nn_results[2].1 == 2);

        // Check distances
        assert_relative_eq!(nn_results[0].2, 0.0);
        assert_relative_eq!(
            nn_results[1].2,
            (0.5_f64.powi(2) + 0.5_f64.powi(2)).sqrt(),
            epsilon = 1e-10
        );
        assert_relative_eq!(nn_results[2].2, 1.0);

        // Test k=0
        let nn_empty = rtree.nearest(&array![0.0, 0.0].view(), 0).unwrap();
        assert_eq!(nn_empty.len(), 0);

        // Test k > size
        let nn_all = rtree.nearest(&array![0.0, 0.0].view(), 20).unwrap();
        assert_eq!(nn_all.len(), 10); // Should return all points
    }

    #[test]
    fn test_rtree_spatial_join() {
        // Create two R-trees
        let mut rtree1: RTree<i32> = RTree::new(2, 2, 4).unwrap();
        let mut rtree2: RTree<char> = RTree::new(2, 2, 4).unwrap();

        // Insert rectangles into the first R-tree
        let rectangles1 = vec![
            (array![0.0, 0.0], array![0.6, 0.6], 0),
            (array![0.4, 0.0], array![1.0, 0.6], 1),
            (array![0.0, 0.4], array![0.6, 1.0], 2),
            (array![0.4, 0.4], array![1.0, 1.0], 3),
        ];

        for (min_corner, max_corner, value) in rectangles1 {
            rtree1
                .insert_rectangle(min_corner, max_corner, value)
                .unwrap();
        }

        // Insert rectangles into the second R-tree
        let rectangles2 = vec![
            (array![0.3, 0.3], array![0.7, 0.7], 'A'),
            (array![0.8, 0.3], array![1.2, 0.7], 'B'),
            (array![0.3, 0.8], array![0.7, 1.2], 'C'),
            (array![0.8, 0.8], array![1.2, 1.2], 'D'),
        ];

        for (min_corner, max_corner, value) in rectangles2 {
            rtree2
                .insert_rectangle(min_corner, max_corner, value)
                .unwrap();
        }

        // Perform a spatial join with an intersection predicate
        let join_results = rtree1
            .spatial_join(&rtree2, |mbr1, mbr2| mbr1.intersects(mbr2))
            .unwrap();

        // There should be multiple pairs since several rectangles intersect
        assert!(
            !join_results.is_empty(),
            "Expected spatial join to find intersecting rectangles"
        );

        // With the given rectangles:
        // Rectangle A [0.3,0.3]x[0.7,0.7] intersects with all 4 rectangles (0,1,2,3)
        // Rectangle B [0.8,0.3]x[1.2,0.7] intersects with rectangles 1 and 3
        // Rectangle C [0.3,0.8]x[0.7,1.2] intersects with rectangles 2 and 3
        // Rectangle D [0.8,0.8]x[1.2,1.2] intersects with rectangle 3
        // Total expected intersections: 4 + 2 + 2 + 1 = 9
        assert_eq!(
            join_results.len(),
            9,
            "Expected 9 intersections, found {}",
            join_results.len()
        );

        // Test a more restrictive join predicate (contains)
        let strict_join_results = rtree1
            .spatial_join(&rtree2, |mbr1, mbr2| mbr1.contains_rectangle(mbr2))
            .unwrap();

        // Should be fewer results than with just intersection
        assert!(strict_join_results.len() <= join_results.len());
    }
}