topn.go

package main

import (
	"container/heap"
	"flag"
	"fmt"
	"math/rand"
	"time"
)

type Region struct {
	ID               int
	StartKey, EndKey []byte
	WriteKeys        int
}

type RegionHeap []*Region

func (h RegionHeap) Len() int           { return len(h) }
func (h RegionHeap) Less(i, j int) bool { return h[i].WriteKeys < h[j].WriteKeys }
func (h RegionHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *RegionHeap) Push(x interface{}) {
	*h = append(*h, x.(*Region))
}

func (h *RegionHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

var r = flag.Int("r", 100000, "region count.")
var n = flag.Int("n", 100, "n.")

func main() {
	flag.Parse()

	regions := make(map[int]*Region)
	for i := 0; i < *r; i++ {
		regions[i] = &Region{
			ID:        i,
			StartKey:  []byte("foobar"),
			EndKey:    []byte("foobar"),
			WriteKeys: rand.Int() % 100000,
		}
	}

	start := time.Now()
	topn := &RegionHeap{}
	for _, r := range regions {
		heap.Push(topn, r)
		if topn.Len() > *n {
			heap.Pop(topn)
		}
	}
	results := make([]*Region, *n)
	for i := *n - 1; i >= 0; i-- {
		results[i] = heap.Pop(topn).(*Region)
	}
	end := time.Now()

	fmt.Println(end.Sub(start))
}