package main
//@description 元素流中获取Top-K的元素,元素需实现Comple接口
//@author chenbintao
//@data 2017-03-31 14:19 初稿
// 2017-03-31 17:44 编译验证通过
import (
"fmt"
"log"
)
//元素接口设计,参与排序的元素,需要实现接口内的方法
type ELEM interface {
//比较当前元素与o元素大小,true:大于o
Comple(interface{}) bool
}
/**
TOP-K-Heap
用于从数据中筛选出前K极值元素
**/
type TopKHeap struct {
heap []ELEM //物理上采用数组实现(接口即指针)
size int64 //堆大小
index int64 //堆元素数(指向下一个可用)
bigMode bool //是否为大顶模式
}
func New(size int64, bigMode bool) *TopKHeap {
if size <= 0 {
return nil
}
//构建一个指定大小的大/小顶堆
this := new(TopKHeap)
this.heap = make([]ELEM, size)
this.size = size
this.index = 0
this.bigMode = bigMode
return this
}
//获取堆容量
func (this *TopKHeap) GetSize() int64 {
return this.size
}
//获取元素数
func (this *TopKHeap) GetCount() int64 {
return this.index
}
//入堆
func (this *TopKHeap) Push(i interface{}) {
if nil == i {
return
}
var o ELEM = i.(ELEM)
if this.index < this.size {
//A.未满,末尾添加,满后建立堆
this.heap[this.index] = o
this.index++
if this.index >= this.size {
this.build()
}
} else {
//B.已满,替换及调整
if ((!o.Comple(this.heap[0])) && this.bigMode) /*入堆元素小于大顶堆的堆顶*/ ||
((o.Comple(this.heap[0])) && !this.bigMode) /*入堆元素大于小顶堆的堆顶*/ {
//满足替换条件
this.heap[0] = o //替换堆顶
this.adjust(0) //堆顶开始调整堆
}
}
return
}
//出堆
func (this *TopKHeap) Pop() (o interface{}) {
if this.index <= 0 ||
this.size <= 0 {
return nil
}
//回溯到尾元素
this.index--
if this.index <= 0 {
this.index = 0
}
o = this.heap[0] //获取堆顶
this.heap[0] = this.heap[this.index] //交换堆顶
this.heap[this.index] = nil //删除尾元
this.adjust(0) //调整堆
return o
}
//建堆
func (this *TopKHeap) build() {
//首个非叶子节点开始调整
i := this.getFather(this.index - 1)
for i >= 0 {
this.adjust(i)
i--
}
return
}
//打印堆
func (this *TopKHeap) String() string {
return fmt.Sprint(this.heap)
}
//调整堆(递归实现)
func (this *TopKHeap) adjust(i int64) int64 {
if this.checkIndex(i) < 0 {
return -1
}
//调整左子堆(交换与递归调整)
var child int64
child = this.getLeft(i)
if child > 0 &&
(((this.heap[i]).Comple(this.heap[child]) && !this.bigMode /*根大于小顶堆孩子*/) ||
(!(this.heap[i]).Comple(this.heap[child]) && this.bigMode /*根小于大顶堆孩子*/)) {
tmp := this.heap[i]
this.heap[i] = this.heap[child]
this.heap[child] = tmp
this.adjust(child)
}
//调整右子堆(交换与递归调整)
child = this.getRight(i)
if child > 0 &&
(((this.heap[i]).Comple(this.heap[child]) && !this.bigMode /*根大于小顶堆孩子*/) ||
(!(this.heap[i]).Comple(this.heap[child]) && this.bigMode /*根小于大顶堆孩子*/)) {
tmp := this.heap[i]
this.heap[i] = this.heap[child]
this.heap[child] = tmp
this.adjust(child)
}
return i
}
//判断堆是否需要调整
func (this *TopKHeap) needAdjust(root, child int64) bool {
if ((this.heap[root]).Comple(this.heap[child]) && !this.bigMode /*根大于小顶堆孩子*/) ||
(!(this.heap[root]).Comple(this.heap[child]) && this.bigMode /*根小于大顶堆孩子*/) {
return true
}
return false
}
//获取父节点(堆顶为0,超出容量时返回-1)
func (this *TopKHeap) getFather(i int64) int64 {
//(i-1)/2向下取整
var j int64
j = (i - 1) / 2
if this.checkIndex(i) >= 0 &&
this.checkIndex(j) >= 0 {
return j
}
return -1
}
//获取左孩子节点(堆顶为0,超出容量时返回-1)
func (this *TopKHeap) getLeft(i int64) int64 {
//(i+1)*2-1
var j int64
j = (i+1)*2 - 1
if this.checkIndex(i) >= 0 &&
this.checkIndex(j) >= 0 {
return j
}
return -1
}
//获取右孩子节点(堆顶为0,超出容量时返回-1)
func (this *TopKHeap) getRight(i int64) int64 {
//(i+1)*2
var j int64
j = (i + 1) * 2
if this.checkIndex(i) >= 0 &&
this.checkIndex(j) >= 0 {
return j
}
return -1
}
//检查节点的合法性(堆顶为0,超出容量时返回-1)
func (this *TopKHeap) checkIndex(i int64) int64 {
if i < 0 ||
i >= this.size ||
i >= this.index ||
nil == this.heap[i] {
return -1
}
return i
}
//===============================================测试程序
//示例用于排序的元素类型
type Elem struct {
i int
}
func (this Elem) Comple(i interface{}) bool {
//ELEM为接口,该方法直接使用Elem类型
var o Elem = i.(Elem)
if this.i > o.i {
return true
} else {
return false
}
}
func (this *Elem) Sring() string {
return fmt.Sprint(this.i)
}
func main() {
heap := New(6, false)
for i := 0; i < 10; i++ {
e := Elem{i: i}
heap.Push(e)
log.Println(heap.String())
}
for heap.GetCount() > 0 {
e := heap.Pop().(Elem)
log.Println(e.Sring())
}
}
运行可参见http://studygolang.com/topics/2602
有疑问加站长微信联系（非本文作者）
