Add problem 973 - K Closest Points To Origin

Author: anirudhology

problems/heaps/k_closest_points_to_origin.go

@@ -0,0 +1,28 @@
+package heaps
+
+import (
+	"container/heap"
+
+	"github.com/anirudhology/leetcode-go/problems/util"
+)
+
+func KClosestPointsToOrigin(points [][]int, k int) [][]int {
+	// List to store k closest points
+	var kClosestPoints [][]int
+	// Special case
+	if len(points) == 0 {
+		return kClosestPoints
+	}
+	minHeap := &util.MinHeapPoint{}
+	heap.Init(minHeap)
+	for i, point := range points {
+		distance := point[0]*point[0] + point[1]*point[1]
+		heap.Push(minHeap, util.Point{Distance: distance, Index: i})
+	}
+	result := make([][]int, k)
+	for i := 0; i < k; i++ {
+		point := heap.Pop(minHeap).(util.Point)
+		result[i] = points[point.Index]
+	}
+	return result
+}

problems/util/min_heap_point.go

@@ -0,0 +1,24 @@
+package util
+
+type Point struct {
+	Distance int
+	Index    int
+}
+
+type MinHeapPoint []Point
+
+func (h MinHeapPoint) Len() int           { return len(h) }
+func (h MinHeapPoint) Less(i, j int) bool { return h[i].Distance < h[j].Distance }
+func (h MinHeapPoint) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *MinHeapPoint) Push(x interface{}) {
+	*h = append(*h, x.(Point))
+}
+
+func (h *MinHeapPoint) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}

tests/heaps/k_closest_points_to_origin_test.go

@@ -0,0 +1,25 @@
+package heaps_test
+
+import (
+	"testing"
+
+	"github.com/anirudhology/leetcode-go/problems/heaps"
+)
+
+func TestKClosest(t *testing.T) {
+	if res := heaps.KClosestPointsToOrigin([][]int{{1, 3}, {-2, 2}}, 1); len(res) != 1 || res[0][0] != -2 || res[0][1] != 2 {
+		t.Errorf("expected [[-2, 2]], got %v", res)
+	}
+	if res := heaps.KClosestPointsToOrigin([][]int{{3, 3}, {5, -1}, {-2, 4}}, 2); len(res) != 2 || res[0][0] != 3 || res[0][1] != 3 || res[1][0] != -2 || res[1][1] != 4 {
+		t.Errorf("expected [[3, 3], [-2, 4]], got %v", res)
+	}
+	if res := heaps.KClosestPointsToOrigin([][]int{{0, 1}, {1, 0}}, 1); len(res) != 1 || res[0][0] != 0 || res[0][1] != 1 {
+		t.Errorf("expected [[0, 1]], got %v", res)
+	}
+	if res := heaps.KClosestPointsToOrigin([][]int{{1, 1}}, 1); len(res) != 1 || res[0][0] != 1 || res[0][1] != 1 {
+		t.Errorf("expected [[1, 1]], got %v", res)
+	}
+	if res := heaps.KClosestPointsToOrigin([][]int{}, 0); len(res) != 0 {
+		t.Errorf("expected [], got %v", res)
+	}
+}