Parallel Optimization of KD-Tree Algorithm: Using C# to Implement Efficient Nearest Neighbor Search

Article Information

Chinese Title: "Parallelized Optimization of KD-Tree Algorithm: Implementing Efficient Nearest Neighbor Search in C#"

English Title: "Parallelized Optimization of KD-Tree Algorithm: Implementing Efficient Nearest Neighbor Search in C#"

Abstract

This article introduces how to use parallel computing techniques to optimize the KD tree algorithm and implements an efficient nearest neighbor search using the C# programming language. First, we briefly introduce the principles and construction process of the KD tree, and then discuss in detail how to leverage parallel computing libraries to build the KD tree in parallel across multiple CPU cores, thereby accelerating the search process. Through experimental verification, we demonstrate that parallel optimization can significantly enhance the construction speed and search efficiency of the KD tree, offering an efficient solution for nearest neighbor search problems in large-scale datasets.

Summary

This article presents a parallelized optimization approach for the KD-tree algorithm and demonstrates an efficient nearest neighbor search implementation in C#. We first introduce the principles and construction process of KD trees, and then discuss in detail how to leverage parallel computing techniques to build KD trees concurrently on multiple CPU cores, thus accelerating the search process. Through experimental validation, we prove that parallelized optimization significantly improves the construction speed and search efficiency of KD trees, providing an efficient solution for nearest neighbor search problems on large-scale datasets.

Version Information

The C# code involved in this article is written using .NET 8.0 and C# 12.

Foreword

Consider the following scenario: there are 1000 locations of Type A (including addresses and GPS coordinates) and 50,000 locations of Type B. The task is to find the nearest Type B location for each Type A location.

My first thought was the KNN algorithm I had previously used in recommendation systems. However, the computational cost of implementing KNN is quite high, and when the data volume is large, it can be very time-consuming. Therefore, for this scenario, I adopted the KD-tree algorithm.

About KD Tree

The KD Tree (K-Dimensional Tree) is a data structure used in multidimensional spaces. It is a variant of the binary tree designed to efficiently organize and search multidimensional data. It is also an efficient algorithm for KNN. The main advantage of a KD Tree is that it can perform fast nearest neighbor and range searches in high-dimensional spaces.

A detailed discussion of the principles behind KD trees is beyond the scope of this article, which provides only a brief introduction. The scenario discussed here concerns the calculation of distances between GPS coordinates, so the dimension chosen for the KD tree is two.

1. Data Structure

• Node: Each node in a KD tree represents a data point. The node contains an axis and a split value, as well as pointers corresponding to the left and right subtrees.

• Root Node: The root node of a KD tree represents the scope of the entire dataset.

• Leaf Node: A leaf node of a KD tree represents a single data point.

2. Construction Process

The construction process of a KD-tree is based on a recursive partitioning strategy. Typically, we select an axis and a split value to divide the dataset into two subsets. Then, we recursively apply the same partitioning operation to each subset until the number of data points in a subset falls below a certain threshold or a specified depth is reached.

During the construction process, the following strategies are generally used to select the axis and split value:

• Axis Selection: The axis is usually determined based on the variance or range of data points along each dimension. The dimension with the largest variance can be chosen as the axis, or dimensions can be selected in a round-robin fashion.

• Split Value Selection: The split value is typically chosen as the median of the data points in the current subset along the selected axis.

3. Search Process

The search process of a KD-tree is also based on recursion. It starts from the root node, traverses the tree downward according to certain rules, until the target data point is found or a leaf node is reached.

During the search process, the following strategies are typically used to determine the search order:

• Determine the splitting direction: Based on the target data point's value on the axis selected at the current node, determine the search direction. If the target value is less than the splitting value of the current node, search the left subtree; otherwise, search the right subtree.

• Determine the search order: Based on the distance of the target data point from the axis selected at the current node, determine the search order. First, search the subtree that is closer, and then search the subtree that is farther away.

• Pruning strategy: During the search process, a pruning strategy can be employed to reduce the number of search branches and improve search efficiency. For example, the distance from the target point to the splitting hyperplane (分割超平面) of the current node can be calculated. If this distance is greater than the current closest distance, that branch can be pruned.

4. Application Scenarios

KD trees are commonly used for organizing and searching data in high-dimensional spaces, especially in the fields of machine learning and data mining. Common applications include nearest neighbor search, range search, and approximate nearest neighbor search.

5. Summary

The KD-tree is an efficient multidimensional spatial data structure suitable for fast nearest-neighbor search and range search. Both its construction and search processes are based on recursive thinking, and search efficiency can be improved by choosing appropriate splitting strategies and pruning strategies.

Distance Calculation

Euclidean Distance

When it comes to distance calculation, the first thing that comes to mind is Euclidean distance, which represents the actual distance between two points in m-dimensional space. The formula is: $$ d = \sqrt{\sum_{i=1}^{k}(x_i - y_i)^2} $$ In two-dimensional and three-dimensional space, Euclidean distance is the distance between two points. Its representation in two-dimensional space is: $$ d = \sqrt{(x_1 - x_2)^2 + (y_1 - y_2)^2} $$ where $(x_1, x_2)$ and $(y_1, y_2)$ are the coordinates of the two points.

When we calculate the distance between two locations here, it seems like using the distance between two points in two-dimensional space would suffice, and that’s what some students have suggested as well.

Distance Between Two Points on a Curved Surface

But then I thought，wait，that’s not right. This only calculates distance on a flat plane，but the Earth isn’t flat—it has curvature.

So I looked up some information and found the Haversine (半正矢) formula.

The name Haversine comes from Ha-VERSINE，which stands for Half-Versine，meaning half of the sine function. The Haversine formula provides a way to calculate the distance between two points on a sphere using their latitudes and longitudes.

$$ haversin(\frac{d}{R}) = haversin(\varphi_2 - \varphi_1) + cos(\varphi_1)cos(\varphi_2)haversin(\Delta\lambda) $$

Where

• d is the distance between two points along the great circle of the sphere (see spherical distance)
• R is the radius of the sphere; for Earth，the average radius can be taken as 6371 km
• φ1, φ2 represent the latitudes of the two points
• Δλ represents the difference in longitudes between the two points

The haversine function hav(θ)，applied to the central angle θ and the differences in latitude and longitude，is defined as $$ hav(\theta) = \sin^2(\frac{\theta}{2}) = \frac{1-\cos \theta}{2} $$ A more detailed derivation of this formula is omitted here. Interested readers are encouraged to consult the references.

Based on the definition of the half-versed sine，the archaversine (反半正弦) can be expressed using the inverse sine: $$ archav(h) = 2 \arcsin \sqrt{h} $$ where $0 \le h \le 1$

Substituting into the Haversine formula，we obtain the formula for solving for the distance d: $$ d = 2R \arcsin (\sqrt{\sin^2(\frac{\varphi_2-\varphi_1}{2}) + \cos \theta_2 \cdot \sin^2(\frac{\lambda_2 - \lambda_1}{2})}) $$

PS: The mathematical consultant for this section: @Wyu-Cnk

Next is the implementation of the Haversine formula in C# to calculate the distance between two points (in kilometers)

public static class DistanceCalculator {
  // Radius of the Earth in kilometers
  private const double EarthRadius = 6371;

  public static double CalculateDistance(ILocation location1, ILocation location2) {
    var lat1 = DegreeToRadian(location1.Lat);
    var lon1 = DegreeToRadian(location1.Lng);
    var lat2 = DegreeToRadian(location2.Lat);
    var lon2 = DegreeToRadian(location2.Lng);

    var dlon = lon2 - lon1;
    var dlat = lat2 - lat1;

    var a = Math.Pow(Math.Sin(dlat / 2), 2) + Math.Cos(lat1) * Math.Cos(lat2) * Math.Pow(Math.Sin(dlon / 2), 2);
    var c = 2 * Math.Atan2(Math.Sqrt(a), Math.Sqrt(1 - a));

    var distance = EarthRadius * c;
    return distance;
  }

  private static double DegreeToRadian(double degree) {
    return degree * Math.PI / 180;
  }
}

Initial Implementation of KD-tree

First, implement a basic KD-tree in C#, including constructing the KD-tree and searching for the nearest node.

Data Structure

To make the program more generic, this document uses interfaces and generics to abstract the data structure to a certain extent.

Data Structure for Locations

public interface ILocation {
  public double Lng { get; set; }
  public double Lat { get; set; }
}

public interface ILocationNode {
  public ILocation Location { get; set; }
}

KD Tree Node (KD树节点)

public class KdTreeNode<T> where T : ILocationNode {
  public T Value { get; set; }
  public KdTreeNode<T>? Left { get; set; }
  public KdTreeNode<T>? Right { get; set; }
}

KDTree Class

A generic parameter is required, which is the same as the generic parameter in the aforementioned KdTreeNode<T> class.

/// <summary>
/// k-dimensional tree
/// </summary>
/// <typeparam name="T"></typeparam>
public class KdTree<T> where T : ILocationNode {
  private readonly IProgress<string> _progress = new Progress<string>(Console.WriteLine);
  private KdTreeNode<T>? _root;
}

KD Tree Construction

Use recursion to construct the KD tree.

public void BuildTree(List<T> items) {
  _root = BuildTree(items, 0, items.Count - 1, 0);
}

private KdTreeNode<T>? BuildTree(List<T> items, int start, int end, int depth) {
  if (start > end)
    return null;

  var axis = depth % 2; // Assuming latitude and longitude as 2D coordinates

  items.Sort(
    (a, b) => axis == 0
    ? a.Location.Lng.CompareTo(b.Location.Lng)
    : a.Location.Lat.CompareTo(b.Location.Lat)
  );

  var medianIndex = start + (end - start) / 2;
  var node = new KdTreeNode<T> {
    Value = items[medianIndex],
    Left = BuildTree(items, start, medianIndex - 1, depth + 1),
    Right = BuildTree(items, medianIndex + 1, end, depth + 1)
  };

  // Progress display
  _progress.Report($"Splitting on depth {depth}, axis {axis}, median index {medianIndex}");

  return node;
}

Usage

Prepare a batch of location data, convert it to the KdNode<T> type, and pass it as a parameter to the BuildTree method for execution.

class Store {
  public string Name { get; set; }
  public string Address { get; set; }
  public Location? Location { get; set; }
}

var stores = new List<Store>(){
  // Example coordinates for New York City
  new Store {
    Name = "Store1",
    Location = new Location { lat = 40.7128, lng = -74.0060 }
  },
  // Example coordinates for Los Angeles
  new Store {
    Name = "Store2",
    Location = new Location { lat = 34.0522, lng = -118.2437 }
  },
  // More data can be loaded from external data sources
};

var kdTree = new KdTree<KdNode<Store>>();
kdTree.BuildTreeParallel(stores.Select(e => new KdNode<Store> {
  Node = e,
  Location = e.Location
}).ToList());

Finding the Nearest Node

The search process of the KD tree is also based on recursion, with the input parameter ILocation representing the location coordinates.

public KdTreeNode<T>? FindNearestNode(ILocation location) {
  return FindNearestNode(_root, location, 0, null);
}

private KdTreeNode<T>? FindNearestNode(KdTreeNode<T>? node, ILocation location, int depth, KdTreeNode<T>? best) {
  if (node == null) return best;

  var bestDistance = best != null
    ? DistanceCalculator.CalculateDistance(location, best.Value.Location)
    : double.PositiveInfinity;
  var currentNodeDistance = DistanceCalculator.CalculateDistance(location, node.Value.Location);

  if (currentNodeDistance < bestDistance)
    best = node;

  var axis = depth % 2;
  var axisDistance =
    axis == 0 ? location.Lng - node.Value.Location.Lng : location.Lat - node.Value.Location.Lat;

  var nearChild = axisDistance < 0 ? node.Left : node.Right;
  var farChild = axisDistance < 0 ? node.Right : node.Left;

  var nearest = FindNearestNode(nearChild, location, depth + 1, best);

  if (nearest != null) {
    var nearestDistance = DistanceCalculator.CalculateDistance(location, nearest.Value.Location);
    if (nearestDistance < bestDistance)
      best = nearest;
  }

  if (Math.Abs(axisDistance) < bestDistance) {
    var farthest = FindNearestNode(farChild, location, depth + 1, best);
    if (farthest != null) {
      var farthestDistance = DistanceCalculator.CalculateDistance(location, farthest.Value.Location);
      if (farthestDistance < bestDistance)
        best = farthest;
    }
  }

  return best;
}

Usage

// Example coordinates for Tokyo City
var location = new Location { lat = 35.652832, lng = 139.839478 }
var nearestNode = kdTree.FindNearestNode(location);

Build Performance

Using C#'s StopWatch tool to time the BuildTree method, with a dataset size of 45066, the above code takes 06:35.9083329 to construct the KD tree, which is 6 minutes and 35 seconds.

Building KD tree...
KD tree construction complete. Time taken: 00:06:35.9083329

Below is the CPU information of the code's runtime environment

12th Gen Intel(R) Core(TM) i7-12700

Base speed:	2.10 GHz
Sockets:	1
Cores:	12
Logical processors:	20
Virtualization:	Enabled
L1 cache:	1.0 MB
L2 cache:	12.0 MB
L3 cache:	25.0 MB

Upon checking the utilization of each CPU core, it was found that only 2 cores were operating at 100% utilization, with an overall load of approximately 20%.

Search Performance

Using FindNearestNode from the code above to find nodes, the execution time is

00:00:00.0242371

Performance Optimization

The current KD tree implementation is already able to meet the requirements of the scenarios described in this article. In terms of performance, the construction process takes a long time, while the search process speed is within an acceptable range. Therefore, the focus of performance optimization is placed on the construction process.

During the construction process, when checking the utilization of each CPU core, it was found that only 2 cores were operating at 100% utilization, with an overall load of around 20%. The reason for this is that, without special optimization, ordinary code runs in a single-threaded manner. Since modern CPUs have a multi-core architecture, a single thread can only utilize a small portion of the CPU's performance. Therefore, my optimization approach is to improve the multi-core utilization of the CPU.

In C#, there are several ways to enable parallel execution in a program. Generally, one would choose methods provided by Parallel or the Task class. Thread can directly manage thread operations and is a lower-level implementation; typically, threads are not managed directly.

Task has a higher level of abstraction. Tasks created using Task (including the async/await statements) are scheduled to run in the thread pool. Task is similar to the concept of "Goroutine" in the Go language (the difference being that Task is implemented through the compiler and state machine at compile time, making it a stackless coroutine, while Goroutine is a stackful coroutine).

Parallel vs Task

Parallel

• The Parallel class provides a simple way to execute parallel loops or iterations. It allows specifying a loop range, automatically splits it into smaller tasks, and then executes these tasks in parallel.
• The primary purpose of the Parallel class is to simplify the writing of parallel loops, enabling developers to easily leverage multi-core processors to improve performance without worrying about the details of managing threads or tasks.
• The Parallel class is typically used to process iterable data collections, such as arrays, lists, etc.

Task

• The Task class offers a more flexible and lower-level mechanism for parallel programming. You can use the Task class to create and manage asynchronous operations, where each Task instance represents an executable unit of work.
• The Task class allows creating tasks with different execution strategies and scheduling options—for example, executing tasks using thread pool threads or new threads, setting task priorities, canceling tasks, and more.
• The Task class is more suited for handling asynchronous operations rather than just parallel loops. It can be used to execute any operation that requires asynchronous execution, such as asynchronous file I/O, network requests, etc.

Summary

• Task offers greater extensibility than Parallel
• Parallel delivers better performance than Task
• Parallel.ForEach is more suitable for CPU-intensive tasks
• Task.WhenAll is more suitable for IO-intensive tasks

For the scenario described in this article, which involves a CPU-intensive task, using Parallel will yield better performance. To compare, this article will implement both Parallel and Task versions.

Optimizing Performance with Task

Refactoring the BuildTree Method

Add a BuildTreeParallel method to KdTree<T>.

public void BuildTreeParallel(List<T> items) {
  Console.WriteLine("Building KD tree...");
  _root = BuildTreeParallel(items, 0, items.Count - 1, 0);
  Console.WriteLine("KD tree construction complete.");
}

private KdTreeNode<T>? BuildTreeParallel(List<T> items, int start, int end, int depth) {
  if (start > end)
    return null;

  var axis = depth % 2; // Assuming latitude and longitude as 2D coordinates

  // 复制一份新的 items 列表
  var sortedItems = items.ToList();

  sortedItems.Sort(
    (a, b) => axis == 0
    ? a.Location.Lng.CompareTo(b.Location.Lng)
    : a.Location.Lat.CompareTo(b.Location.Lat)
  );

  var medianIndex = start + (end - start) / 2;
  var node = new KdTreeNode<T> {
    Value = sortedItems[medianIndex]
  };

  var leftTask = Task.Run(() => BuildTreeParallel(sortedItems, start, medianIndex - 1, depth + 1));
  var rightTask = Task.Run(() => BuildTreeParallel(sortedItems, medianIndex + 1, end, depth + 1));

  node.Left = leftTask.Result;
  node.Right = rightTask.Result;

  return node;
}

The code logic is basically the same as the single-threaded version. The difference is that in the parallel version, the leaf nodes of each node are constructed within a new Task.

Also, regarding the sorting part of items, in the multithreaded version, sorting cannot be performed directly. Instead, a new copy of the list must be made for sorting. In high-concurrency scenarios, sorting the same items object can lead to consistency errors. The specific error appears as:

System.AggregateException: One or more errors occurred. (Unable to sort because the IComparer.Compare() method returns inconsistent result
s. Either a value does not compare equal to itself, or one value repeatedly compared to another value yields different results. IComparer:
'System.Comparison`1[StoreProximityBroadbandLocator.Services.KdNode`1[StoreProximityBroadbandLocator.Models.Device]]'.)

Essentially, this means that the comparator's (IComparer) comparison method returned inconsistent results. Making a new copy of the items list for sorting resolves this issue.

Refactoring the FindNearestNode Method

In fact, searching for nodes involves minimal computation, and parallel optimization offers little benefit. Instead, it can introduce additional overhead due to thread scheduling.

This article provides a Task-based implementation, but does not recommend using this approach.

public KdTreeNode<T>? FindNearestNodeParallel(ILocation location) {
  return FindNearestNodeParallel(_root, location, 0, null);
}

private KdTreeNode<T>? FindNearestNodeParallel(KdTreeNode<T>? node, ILocation location, int depth, KdTreeNode<T>? best) {
  if (node == null) return best;

  var bestDistance = best != null
    ? DistanceCalculator.CalculateDistance(location, best.Value.Location)
    : double.PositiveInfinity;
  var currentNodeDistance = DistanceCalculator.CalculateDistance(location, node.Value.Location);

  if (currentNodeDistance < bestDistance)
    best = node;

  var axis = depth % 2;
  var axisDistance =
    axis == 0 ? location.Lng - node.Value.Location.Lng : location.Lat - node.Value.Location.Lat;

  var nearChild = axisDistance < 0 ? node.Left : node.Right;
  var farChild = axisDistance < 0 ? node.Right : node.Left;

  var nearestTask = Task.Run(() => FindNearestNodeParallel(nearChild, location, depth + 1, best));
  var nearest = nearestTask.Result;

  if (nearest != null) {
    var nearestDistance = DistanceCalculator.CalculateDistance(location, nearest.Value.Location);
    if (nearestDistance < bestDistance)
      best = nearest;
  }

  if (Math.Abs(axisDistance) < bestDistance) {
    var farthestTask = Task.Run(() => FindNearestNodeParallel(farChild, location, depth + 1, best));
    var farthest = farthestTask.Result;
    if (farthest != null) {
      var farthestDistance = DistanceCalculator.CalculateDistance(location, farthest.Value.Location);
      if (farthestDistance < bestDistance)
        best = farthest;
    }
  }

  return best;
}

Performance Testing

Similar to the single-threaded version, run the KD tree construction and node lookup methods.

Now it can fully utilize the entire CPU. The development environment for this article uses an Intel i7-12700 CPU, with an all-core frequency reaching up to approximately 3.6GHz.

The BuildTree method, implemented based on Task, takes about 48 seconds to execute.

Performance Optimization with Parallel

Refactoring the BuildTree Method

Add the BuildTreeParallel2 method to KdTree<T>.

public void BuildTreeParallel2(List<T> items) {
  _root = new KdTreeNode<T>();
  BuildTreeParallel2(_root, items, 0, items.Count - 1, 0);
  Console.WriteLine("KD tree construction complete.");
}

private void BuildTreeParallel2(KdTreeNode<T> node, List<T> items, int start, int end, int depth) {
  if (start > end)
    return;

  var axis = depth % 2; // Assuming latitude and longitude as 2D coordinates

  // Create a copy of the items list
  var sortedItems = items.ToList();

  sortedItems.Sort(
    (a, b) => axis == 0
    ? a.Location.Lng.CompareTo(b.Location.Lng)
    : a.Location.Lat.CompareTo(b.Location.Lat)
  );

  var medianIndex = start + (end - start) / 2;
  node.Value = sortedItems[medianIndex];

  if (start < medianIndex) {
    node.Left = new KdTreeNode<T>();
  }

  if (medianIndex < end) {
    node.Right = new KdTreeNode<T>();
  }

  Parallel.Invoke(
    () => {
      if (node.Left != null) {
        BuildTreeParallel2(node.Left, sortedItems, start, medianIndex - 1, depth + 1);
      }
    },
    () => {
      if (node.Right != null) {
        BuildTreeParallel2(node.Right, sortedItems, medianIndex + 1, end, depth + 1);
      }
    }
  );
}

Error Example

At the beginning, I wrote some incorrect code.

The build process worked fine, but every time I tried to search, I would encounter a null reference exception.

After debugging, I finally found the cause of the error.

Here is the original code attached, and then let's analyze it.

node.Left = new KdTreeNode<T>();
node.Right = new KdTreeNode<T>();

Parallel.Invoke(
  () => BuildTreeParallel2(node.Left, sortedItems, start, medianIndex - 1, depth + 1),
  () => BuildTreeParallel2(node.Right, sortedItems, medianIndex + 1, end, depth + 1)
);

The old version of the code is more concise.

It seems fine at first glance, but because there was initially a line if (start > end) return;, after a certain branch was already constructed, the end nodes of that branch would still generate an empty node.

It's not entirely empty—it's a KdTreeNode<T> object, but its Value and Location properties are all null.

Therefore, I had to change it to the method shown above, which not only fully utilizes the CPU but also resolves this issue.

Performance Testing

Using C#'s StopWatch tool for testing, the BuildTree method implemented with Parallel took approximately 43 seconds to execute, which is indeed faster compared to the Task implementation.

References

• https://en.wikipedia.org/wiki/Haversine_formula
• A Rabbit Helps You Understand kNN
• The Concept of kd-tree Algorithm
• Detailed Explanation of the kd-tree Algorithm - https://zhuanlan.zhihu.com/p/23966698
• In-depth Explanation of KD-Tree Principles - https://zhuanlan.zhihu.com/p/112246942
• A Guide to Parallel Execution in C# : Harness the Power of the Task Parallel Library (TPL)
• From C# to golang: less is more
• When would you use Parallel.ForEachAsync() and when Task.WhenAll()
• Task.WhenAll vs Parallel.Foreach