98 lines
1.5 KiB
Go
98 lines
1.5 KiB
Go
package main
|
|
|
|
import (
|
|
"container/heap"
|
|
"fmt"
|
|
"math/rand"
|
|
"sort"
|
|
"time"
|
|
)
|
|
|
|
func main() {
|
|
for i := range merge(source(500), source(100), source(200)) {
|
|
fmt.Printf("%20d\n", i)
|
|
}
|
|
}
|
|
|
|
func source(c int) <-chan int {
|
|
out := make(chan int)
|
|
go func() {
|
|
vals := make([]int, c)
|
|
src := rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
for i := 0; i < c; i++ {
|
|
vals[i] = src.Int()
|
|
}
|
|
|
|
sort.Ints(vals)
|
|
|
|
for _, i := range vals {
|
|
out <- i
|
|
}
|
|
close(out)
|
|
}()
|
|
return out
|
|
}
|
|
|
|
func merge(cs ...<-chan int) <-chan int {
|
|
out := make(chan int)
|
|
|
|
go func() {
|
|
h := &IntHeap{}
|
|
heap.Init(h)
|
|
|
|
// prime the pumps
|
|
for i, src := range cs {
|
|
head, ok := <-src
|
|
if !ok {
|
|
continue
|
|
}
|
|
it := item{
|
|
val: head,
|
|
src: i,
|
|
}
|
|
heap.Push(h, it)
|
|
}
|
|
|
|
for h.Len() > 0 {
|
|
top := heap.Pop(h).(item)
|
|
out <- top.val
|
|
if head, ok := <-cs[top.src]; ok {
|
|
it := item{
|
|
val: head,
|
|
src: top.src,
|
|
}
|
|
heap.Push(h, it)
|
|
}
|
|
}
|
|
|
|
close(out)
|
|
}()
|
|
return out
|
|
}
|
|
|
|
// An IntHeap is a min-heap of ints.
|
|
type IntHeap []item
|
|
|
|
func (h IntHeap) Len() int { return len(h) }
|
|
func (h IntHeap) Less(i, j int) bool { return h[i].val < h[j].val }
|
|
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
|
|
|
|
func (h *IntHeap) Push(x interface{}) {
|
|
// Push and Pop use pointer receivers because they modify the slice's length,
|
|
// not just its contents.
|
|
*h = append(*h, x.(item))
|
|
}
|
|
|
|
func (h *IntHeap) Pop() interface{} {
|
|
old := *h
|
|
n := len(old)
|
|
x := old[n-1]
|
|
*h = old[0 : n-1]
|
|
return x
|
|
}
|
|
|
|
type item struct {
|
|
val int
|
|
src int
|
|
}
|