Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save alldroll/8decefbbb043d5f185375dc1d74aa158 to your computer and use it in GitHub Desktop.
Save alldroll/8decefbbb043d5f185375dc1d74aa158 to your computer and use it in GitHub Desktop.
import "container/heap"
type HeapItem struct {
value int
listID int
}
type Heap []HeapItem
func (h Heap) Len() int { return len(h) }
func (h Heap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h Heap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *Heap) Push(x interface{}) { *h = append(*h, x.(HeapItem)) }
func (h *Heap) Pop() interface{} {
n := h.Len()
old := *h
x := old[n - 1]
*h = old[:n - 1]
return x
}
func smallestRange(nums [][]int) []int {
k := len(nums)
r := []int{(1 << 31) - 1, -1 << 31}
h := Heap{}
for i, list := range nums {
heap.Push(&h, HeapItem{
value: list[0],
listID: i,
})
r[0] = min(r[0], list[0])
r[1] = max(r[1], list[0])
}
result := []int{r[0], r[1]}
for h.Len() == k {
item := h[0]
nums[item.listID] = nums[item.listID][1:]
if (r[1] - r[0]) < (result[1] - result[0]) {
result[0], result[1] = r[0], r[1]
}
if len(nums[item.listID]) == 0 {
break
}
h[0].value = nums[item.listID][0]
heap.Fix(&h, 0)
r[0] = h[0].value
r[1] = max(r[1], nums[item.listID][0])
}
return result
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
func max(a, b int) int {
if a < b {
return b
}
return a
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment