Compact multidimensional sorting and searching with kdtools

Introduction

Sorting and searching are fundamental operations in computer and data science. The objective of the kdtools is to expose a C++ header file implementing efficient sorting and searching methods for multidimensional tuple-like data. The interface mirrors that of the C++ Standard Template Library (STL) sorting and searching functions. It is a single header-only library with no external dependencies other than std.

A common approach to improving search performance is to place data into an ordered tree structure, the most common being a binary tree and its variants. Computer textbooks and introductory courses are littered with examples of tree structures, however it has been pointed out by practitioners that many tree-type data structures have considerable storage overhead and poor memory locality. This results from using linked-list-like pointers to connect nodes. In many cases, a fully or partially sorted sequence combined with binary search will outperform pointer-based trees. Indeed, the STL contains a set of algorithms for sorting or partitioning a data sequence and searching within partitioned sequences.

A limitation of the STL algorithms and data structures is that they can only be ordered in a single dimension. Sort keys can possess higher cardinality however comparison operators and ordering are unidimensional. This is suboptimal for data that are intrinsically multidimensional, such as geographical coordinates or other multivariate data. In multiple dimensions, it is preferred to order a sequence of keys such the average distance between neighboring pairs in the sequence is small. There exist a vast array of techniques to accomplish sequence localization; however I will only discuss one known as the kd-tree.

A kd-tree is a type of binary tree that cycles among the dimensions of its stored objects at each level of the tree hierarchy. In the case of spatial data, the root node will contain the median key in the x-direction and each of its daughter nodes will contain the median in the y-direction of their respective partitions. Subsequent daughter nodes will again contain medians in the x-direction, and so on until there are no additional keys to the left or right. Searching a kd-tree involves recursive comparisons, cycling among the dimensions, until a leaf-node is encountered. Most searching operations can be accomplished in logarithmic time complexity yielding an efficient method.

Precisely as an ordinary binary tree can be replaced by a sorted range, a sequence of multidimensional elements can be sorted, in-place if desired, via recursive partitioning following the kd-tree scheme. Divide-and-conquer recursive partitioning is a well-known sorting algorithm and is the basis of the quicksort algorithm. Partitioning a sequence places all elements less than (or more generally meeting some predicate) to the left of a selected key, while all elements greater-than-or-equal-to follow to the right. The kd-sort algorithm presented here simply modifies quicksort to cycle among the dimensions of the stored elements. Searching follows exactly as the kd-tree except that nodes correspond to pivot-elements in the linear array.

Implementation of kd_sort enables a set of additional algorithms for searching and nearest neighbor queries. The following table gives the exposed functions and corresponding STL routines.

Sorting

The kd-sort algorithm in the kdtools package is implemented as a C++ function template parameterized by the starting dimension (usually 0) and the iterator type, which will be auto-deduced by the compiler.

template <size_t I, typename Iter>
void kd_sort(Iter first, Iter last)
{
  using TupleType = iter_value_t<Iter>;
  constexpr auto J = next_dim<I, TupleType>::value;
  if (distance(first, last) > 1)
  {
    auto pred = kd_less<I>();
    auto pivot = middle_of(first, last);
    nth_element(first, pivot, last, pred);
    pivot = adjust_pivot(first, pivot, pred);
    kd_sort<J>(next(pivot), last);
    kd_sort<J>(first, pivot);
  }
}

The code utilizes the nth_element function from the STL. The nth_element algorithm partially sorts a range such that the element pointed to by the pivot iterator is the one that would be in that position if the range was fully sorted. Since the pivot is chosen as the middle element, after calling nth_element, it points to the median key. The nth_element algorithm furthermore partitions the data such that all of the elements to the right of the pivot are not-less-than elements to the left. This is sufficient to recursively sort the range, however it will fail when there are enough ties that the first instance of the median key falls to the left of the pivot. A second partitioning pass, adjust_pivot is therefore required to ensure the pivot points to the first occurrence of the median key in the sequence. Adjust_pivot is thin wrapper around the partition algorithm from the STL.

Once the first dimension is partitioned, the ranges right and left of the pivot are partitioned, and so on until first and last span a single or no value. The dimension index I is incremented at compile-time using template metaprogramming. The value member of the templated struct next_dim holds the incremented value of I that has been wrapped to cycle through successive dimensions of TupleType. The result of kd_sort is a sequence ordered as if it were inserted into a kd-tree data structure. The following figure illustrates the resulting order in 2-dimensions.

Here the sequence order is shown with a continuous color gradient. The resulting patchiness indicates a hierarchical tree-like partitioning of the space providing the desired locality. Placing nearby points in proximity is what allows for efficient searches once the data are sorted. Sort times for vectors of tuple-like objects are relatively efficient, taking roughly twice the time of an ordinary lexicographical sort.

In the above results, lex_sort is the STL built-in std::sort algorithm using kd_less as the predicate. The results indicate that kd_sort is reasonably fast, requiring roughly twice the time as the STL built-in sorting algorithm. Because of the divide-and-conquer nature of kd_sort, it is trivially parallelizable. The threaded version is equal to or faster than the regular sort. Performance comparable to the build-in STL sort function indicates that this reference implementation is reasonable. The computational complexity appears to scale linearly in the number of tuples in all cases.

template <size_t I, size_t K = 0>
struct kd_less
{
  template <typename TupleType>
  typename enable_if<is_not_last<K, TupleType>::value, bool>::type
  operator()(const TupleType& lhs, const TupleType& rhs) const
  {
    constexpr auto J = next_dim<I, TupleType>::value;
    return get<I>(lhs) == get<I>(rhs) ?
      kd_less<J, K + 1>()(lhs, rhs) :
        get<I>(lhs) < get<I>(rhs);
  }
  template <typename TupleType>
  typename enable_if<is_last<K, TupleType>::value, bool>::type
  operator()(const TupleType& lhs, const TupleType& rhs) const
  {
    return get<I>(lhs) < get<I>(rhs);
  }
};

j = cbind(sample(1:9))
k = cbind(sample(1:9))
t(kd_sort(j))

t(kd_sort(cbind(j, k)))

t(kd_sort(cbind(0, j, k)))

t(kd_sort(cbind(j, k, 0)))

t(kd_sort(cbind(0, j, 0)))

template <size_t I,
          typename Iter,
          typename TupleType,
          typename OutIter>
void kd_range_query(Iter first, Iter last,
                    const TupleType& lower,
                    const TupleType& upper,
                    OutIter outp)
{
  if (distance(first, last) > 32) {
    auto pred = less_nth<I>();
    auto pivot = find_pivot<I>(first, last);
    constexpr auto J = next_dim<I, TupleType>::value;
    if (within(*pivot, lower, upper)) *outp++ = *pivot;
    if (!pred(*pivot, lower)) // search left
      kd_range_query<J>(first, pivot, lower, upper, outp);
    if (pred(*pivot, upper)) // search right
      kd_range_query<J>(next(pivot), last, lower, upper, outp);
  } else {
    copy_if(first, last, outp, [&](const TupleType& x){
      return within(x, lower, upper);
    });
  }
  return;
}

// [[Rcpp::depends(kdtools)]]
// [[Rcpp::plugins(cpp11)]]

#include "mixed_query.h"

// [[Rcpp::export]]
void mixed_query(NumericVector c1, CharacterVector c2)
{
  auto n = c1.size();
  
  range_type data;
  data.reserve(n);
  
  transform(begin(c1), end(c1), begin(c2), back_inserter(data), make_key());
  
  auto t = time_it([&]{
    
    kd_sort(begin(data), end(data));
    
  });
  Rcout << "\nSort time: " << t.count() << " seconds" << std::endl;
  
  Rcout << "\nFirst and last 5 elements of sorted data:\n" << std::endl;
  print_range(data);
  
  range_type result;
  
  auto u = time_it([&]{
    
    kd_range_query(begin(data), end(data),
                   key_type(0.4, "w"), key_type(0.6, "z"),
                   back_inserter(result));
    
  });
  
  Rcout << "\nQuery time: " << u.count() << " seconds" << std::endl;

  Rcout << "\nFirst and last 5 elements of query return:\n" << std::endl;
  kd_sort(begin(result), end(result)); print_range(result);
  
  result.clear();
  
  auto v = time_it([&]{
    
    kd_nearest_neighbors(begin(data), end(data),
                         key_type(0.5, "kdtools"), 100,
                         back_inserter(result));
    
  });
  
  Rcout << "\nNeighbors query time: " << v.count() << " seconds" << std::endl;

  Rcout << "\nFirst and last 5 elements of neighbors query return:\n" << std::endl;
  print_range(result);
}

data("parts_of_speech")
numbers = runif(nrow(parts_of_speech))
strings = sample(tolower(parts_of_speech[[1]]))
mixed_query(signif(numbers, 4), strings)
#> 
#> Sort time: 0.137954 seconds
#> 
#> First and last 5 elements of sorted data:
#> 
#> 0 0.0001374 about-ship
#> 1 0.00196 absorptiometer
#> 2 0.003099 absorb
#> 3 0.002444 academical
#> 4 0.00174 acanthite
#> ...
#> 208254 0.9954 yorktown
#> 208255 0.9986 zach
#> 208256 0.9995 zecchini
#> 208257 0.9952 ziagos
#> 208258 0.9954 zetes
#> 
#> Query time: 0.000138909 seconds
#> 
#> First and last 5 elements of query return:
#> 
#> 0 0.4058 waggon
#> 1 0.4035 waggon-headed
#> 2 0.4101 wagga
#> 3 0.4116 wahlstrom
#> 4 0.4095 walrus
#> ...
#> 1078 0.5967 xenophontean
#> 1079 0.596 xmas
#> 1080 0.5837 yup
#> 1081 0.5939 yenta
#> 1082 0.597 youthen
#> 
#> Neighbors query time: 0.00103432 seconds
#> 
#> First and last 5 elements of neighbors query return:
#> 
#> 0 0.2749 lota
#> 1 0.06498 kamloops
#> 2 0.7165 lose
#> 3 0.2892 kroo
#> 4 0.6664 kingwood
#> ...
#> 95 0.4529 kelotomy
#> 96 0.4709 kelso
#> 97 0.6921 katholikos
#> 98 0.3054 koodoos
#> 99 0.4208 doltish

template<class InputIt1, class InputIt2>
double set_similarity(InputIt1 first1, InputIt1 last1,
                      InputIt2 first2, InputIt2 last2)
{
  double num = 0, denom = 0;
  while (first1 != last1 && first2 != last2) {
    if (*first1 < *first2) {
      ++first1; ++denom;
    } else {
      if (!(*first2 < *first1)) {
        ++first1; ++num;
      }
      ++first2; ++denom;
    }
  }
  denom += std::distance(first1, last1) + std::distance(first2, last2);
  return num / denom;
}

namespace kdtools {

template <>
double scalar_diff(const std::string& lhs, const std::string& rhs)
{
  std::string a(lhs), b(rhs);
  std::sort(begin(a), end(a));
  std::sort(begin(b), end(b));
  return 1 - set_similarity(begin(a), end(a),
                            begin(b), end(b));
}

}; // namespace kdtools

Conclusions

The kdtools package implements efficient sorting and searching of arbitary tuple-like object in C++. One of the major design decisions was to make the number of dimensions fixed at compile time. This is a reasonable tradeoff as these methods break down in high dimensions where search times will be no better than that of brute-force random searching. Nonetheless, a runtime-dimensioned extension would be easy to add and could allow working with runtime-determined types as is common in R. The current implementation is a demonstration, however it does illustrate the value of the STL as the algorithms are constructed from STL components. The generic nature of the kdtools implementation is an attractive feature as it permits flexibility and can be adapted to arbitrary tuple-like data as needed. Modern development of C++ has emphasized tuple-style data and algorithms and kdtools extends these developments.