堆实现的Top-K算法,元素流中筛选极值集合-博文

hansegod · · 1407 次点击 · · 开始浏览    
这是一个创建于 的文章,其中的信息可能已经有所发展或是发生改变。

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


有疑问加站长微信联系(非本文作者)

入群交流(和以上内容无关):加入Go大咖交流群,或添加微信:liuxiaoyan-s 备注:入群;或加QQ群:692541889

1407 次点击  
加入收藏 微博
暂无回复
添加一条新回复 (您需要 登录 后才能回复 没有账号 ?)
  • 请尽量让自己的回复能够对别人有帮助
  • 支持 Markdown 格式, **粗体**、~~删除线~~、`单行代码`
  • 支持 @ 本站用户;支持表情(输入 : 提示),见 Emoji cheat sheet
  • 图片支持拖拽、截图粘贴等方式上传