Merge pull request #627 from vaskoz/day306

Day306

README.md

@@ -314,3 +314,4 @@ problems from
 * [Day 303](https://github.com/vaskoz/dailycodingproblem-go/issues/616)
 * [Day 304](https://github.com/vaskoz/dailycodingproblem-go/issues/618)
 * [Day 305](https://github.com/vaskoz/dailycodingproblem-go/issues/622)
+* [Day 306](https://github.com/vaskoz/dailycodingproblem-go/issues/626)

day306/problem.go

@@ -0,0 +1,43 @@
+package day306
+
+import "container/heap"
+
+type intHeap []int
+
+func (h intHeap) Len() int           { return len(h) }
+func (h intHeap) Less(i, j int) bool { return h[i] < h[j] }
+func (h intHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *intHeap) Push(x interface{}) {
+	*h = append(*h, x.(int))
+}
+
+func (h *intHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}
+
+// SortK sorts the input that is shifted by at most
+// k positions and returns a new output.
+// Runs in O(N log K) time.
+func SortK(nums []int, k int) []int {
+	size := 2*k + 1
+	h := make(intHeap, 0, size)
+	heap.Init(&h)
+	sorted := make([]int, 0, len(nums))
+	for i, num := range nums {
+		if i >= size {
+			smallest := heap.Pop(&h).(int)
+			sorted = append(sorted, smallest)
+		}
+		heap.Push(&h, num)
+	}
+	for h.Len() > 0 {
+		smallest := heap.Pop(&h).(int)
+		sorted = append(sorted, smallest)
+	}
+	return sorted
+}

day306/problem_test.go

@@ -0,0 +1,34 @@
+package day306
+
+import (
+	"reflect"
+	"testing"
+)
+
+// nolint
+var testcases = []struct {
+	nums   []int
+	k      int
+	sorted []int
+}{
+	{[]int{3, 2, 1, 6, 5, 4, 9, 8, 7}, 1, []int{1, 2, 3, 4, 5, 6, 7, 8, 9}},
+	{[]int{6, 5, 3, 2, 8, 10, 9}, 3, []int{2, 3, 5, 6, 8, 9, 10}},
+	{[]int{10, 9, 8, 7, 4, 70, 60, 50}, 4, []int{4, 7, 8, 9, 10, 50, 60, 70}},
+}
+
+func TestSortK(t *testing.T) {
+	t.Parallel()
+	for _, tc := range testcases {
+		if sorted := SortK(tc.nums, tc.k); !reflect.DeepEqual(sorted, tc.sorted) {
+			t.Errorf("Expected %v, got %v", tc.sorted, sorted)
+		}
+	}
+}
+
+func BenchmarkSortK(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		for _, tc := range testcases {
+			SortK(tc.nums, tc.k)
+		}
+	}
+}