Coding Interview Questions | K Closest Points to Origin

January 8th, 2024

K Closest Points to Origin Problem Introduction:

This summary will talk about my solution to the K closest points to the origin problem as seen on leetcode here.

How do you solve the K Closest Points to the Origin Problem in O(n) time?:

Solution:

type PointDistance = {
  distance: number;
  idx: number;
};

class MinHeap {
  private heap: PointDistance[] = [];

  private heapify(): void {
    const n = this.heap.length;
    // Start from the last non-leaf node and move up
    for (let i = Math.floor(n / 2) - 1; i >= 0; i--) {
      this.heapifyDown(i);
    }
  }

  buildHeap(values: PointDistance[]): void {
    this.heap = values;
    this.heapify();
  }

  private swap(index1: number, index2: number): void {
    [this.heap[index1], this.heap[index2]] = [
      this.heap[index2],
      this.heap[index1],
    ];
  }

  private getParentIndex(index: number): number {
    return Math.floor((index - 1) / 2);
  }

  private getLeftChildIndex(index: number): number {
    return 2 * index + 1;
  }

  private getRightChildIndex(index: number): number {
    return 2 * index + 2;
  }

  private heapifyUp(): void {
    let currentIndex = this.heap.length - 1;

    while (currentIndex > 0) {
      const parentIndex = this.getParentIndex(currentIndex);

      if (this.heap[parentIndex].distance > this.heap[currentIndex].distance) {
        this.swap(parentIndex, currentIndex);
        currentIndex = parentIndex;
      } else {
        break;
      }
    }
  }

  private heapifyDown(idx?: number): void {
    let currentIndex = idx ? idx : 0;

    while (true) {
      const leftChildIndex = this.getLeftChildIndex(currentIndex);
      const rightChildIndex = this.getRightChildIndex(currentIndex);
      let smallestChildIndex = null;

      if (
        leftChildIndex < this.heap.length &&
        this.heap[leftChildIndex].distance < this.heap[currentIndex].distance
      ) {
        smallestChildIndex = leftChildIndex;
      }

      if (
        rightChildIndex < this.heap.length &&
        this.heap[rightChildIndex].distance < this.heap[currentIndex].distance
      ) {
        smallestChildIndex =
          this.heap[rightChildIndex].distance <
          this.heap[leftChildIndex].distance
            ? rightChildIndex
            : smallestChildIndex;
      }

      if (smallestChildIndex === null) {
        break;
      }

      this.swap(currentIndex, smallestChildIndex);
      currentIndex = smallestChildIndex;
    }
  }

  insert(value: PointDistance): void {
    this.heap.push(value);
    this.heapifyUp();
  }

  extractMin(): PointDistance | undefined {
    if (this.heap.length === 0) {
      return undefined;
    }

    if (this.heap.length === 1) {
      return this.heap.pop()!;
    }

    const min = this.heap[0];
    this.heap[0] = this.heap.pop()!;
    this.heapifyDown();

    return min;
  }

  peekMin(): PointDistance | undefined {
    return this.heap.length > 0 ? this.heap[0] : undefined;
  }

  size(): number {
    return this.heap.length;
  }

  isEmpty(): boolean {
    return this.heap.length === 0;
  }
}

function kClosest(points: number[][], k: number): number[][] {
  // Loop through once, calculating distance to origin for each point using c = sqrt(a2 + b2), populating a distances array
  // [[1, 1], [1,2]] -> [s2, s5]
  // create min heap?

  const distances = points.map((point, idx) => {
    return { distance: point[0] ** 2 + point[1] ** 2, idx };
  });
  const minDistancesHeap = new MinHeap();
  minDistancesHeap.buildHeap(distances);

  const kClosestPoints: number[][] = [];

  for (let i = 0; i < k; i++) {
    const currentMin = minDistancesHeap.extractMin()!;

    kClosestPoints.push(points[currentMin.idx]);
  }
  return kClosestPoints;
}

K Closest Points Solution Summary:

Below is a breakdown of the key aspects of the solution above:

1. Min Heap Construction: The MinHeap class is implemented with methods for heap construction, including heapifyUp, heapifyDown, and buildHeap to efficiently manage a heap of PointDistance objects.

2. Distance Calculation: The solution calculates the distance to the origin for each point using the Euclidean distance formula sqrt(a^2 + b^2) and populates a distances array with {distance, idx} objects.

3. Min Heap Usage: The distances array is used to construct a min heap using the MinHeap class, ensuring quick retrieval of the k closest points.

4. K Closest Points Extraction: The algorithm then iteratively extracts the minimum distance point from the min heap and appends the corresponding point from the original array to the result array kClosestPoints.

Complexities

1. Time Complexity: The time complexity of the solution is O(n + k log n), where n is the number of points. Constructing the distances array takes O(n) time, and building the min heap and extracting the k closest points takes O(k log n) time. If k << n, the overall time complexity becomes O(n).
2. Space Complexity: The space complexity is O(n) due to the distances array and the min heap.