scala写算法-用小根堆解决topK

Posted 2020-10-18 蒋航的博客

topK问题是指从大量数据中获取最大(或最小)的k个数,比如从全校学生中寻找成绩最高的500名学生等等.
本问题可采用小根堆解决.思路是先把源数据中的前k个数放入堆中,然后构建堆,使其保持堆序(可以简单的看成k次insert操作).然后从源数据中的第k个数据之后的每个元素与堆的根节点(小根堆得root是最小的)比较,如果小于root,那么直接pass;如果大于,则执行headp.deleteMin,然后把该元素插入堆中并再次保持堆序.保持堆序需要涉及上滤与下滤的过程.
样例为:

object Main extends App{
  val array=Array(0,14,16,19,3,3768,345,3,343,545,455,7567,657,67,65756,756,756,756,7657,657,657,4,534,535,345343,423,4,46546,546,544,546,546,345,345,435,34534,53,5345,45,45,435,43,54,35,435,435,34,5,435,45,65,65,6576,7,65,56,7,45,43,543,53,453,45,345,34)
  //数组从1开始
  val k=3
  val heap=new Heap
  1 to k foreach(i=>heap.insert(array(i)))
  k+1 to array.length-1 foreach(i=>{
    while (heap.getSize>10){
      heap.deleteMin
    }
    if(array(i)>heap.getMin){
      heap.deleteMin
      heap.insert(array(i))
    }
  })
  heap.print
}

再来说说,Heap是如何实现的
运用scala常常遇到这样的选择:变量的变与不变,这是一个问题.
还有for推导式,实际上是map flatMap等的语法糖.

代码:

class Heap {
  type T=Int
  private var size=0
  val elemnts=new Array[T](200) // 把index=0处的位置空出 2n 2n+1 n/2
  def getSize:Int=return size
  def insert(x:T):Unit={
    def loop(i:Int):Int={
      if(elemnts(i/2)<=x)
        return i
      elemnts(i)=elemnts(i/2)
      loop(i/2)
    }
    val i=loop(size+1)
    elemnts(i)=x
    size+=1
  }
  def deleteMin:T={
    def loop(i:Int):Int={
      if(i>size) return i/2
      val left=elemnts(2*i)
      val right=elemnts(2*i+1)
      if(left<right){
        elemnts(i)=left
        loop(i*2)
      }else{
        elemnts(i)=right
        loop(2*i+1)
      }
    }
    val result=elemnts(1)
    val last=elemnts(size)
    val i=loop(1)
    elemnts(i)=last
    size-=1
    return result
  }
  def print:Unit= 1 to size foreach(i=>println(elemnts(i)))
  def getMin:T=return elemnts(1)
  //代码亲测无误
}

在实现insert与deleteMin时,由于scala并没有break关键字(虽然你可以使用Breakable这个类实现,实际上通过抛出异常模拟break,不灵活),为实现上虑(insert),考虑用递归来模拟for循环.
代码:

  def insert(x:T):Unit={
    def loop(i:Int):Int={
      if(elemnts(i/2)<=x)
        return i //如果父亲节点比待插入值x小,则本节点应该插入x
      elemnts(i)=elemnts(i/2) //上虑
      loop(i/2)
    }
    val i=loop(size+1) //返回待x插入的位置
    elemnts(i)=x
    size+=1
  }

相比从c语言版,基于scala的代码还是容易记忆与相当稳健:

void insert(Element x,Heap* h){
    int i=0;
    if(isFull(h))
        error("full");
    for(i=++h->size;h->elements[i/2]>x;i/=2)
        h->elements[i]=h->elements[i/2];
    h->elements[i]=x;
}