golang之TopN算法

朱雀 2022-06-13 07:26 147阅读 0赞

相信大家已经了解了TOPN算法的原理了  
不多说 直接上代码，其中\[\]leveldb.UserExp 看需求来定义

//h[root] parentNode h[root/2] childNode h[root*2+1] h[root*2+2]
    func GetTopN(n int, a []leveldb.UserExp) ([]leveldb.UserExp, error) {
        h := make([]leveldb.UserExp, 0)
        lenA := len(a)
        if n >= lenA {
            return h, errors.New("the top n greater than the length of a")
        }
        for i := 0; i < lenA; i++ {
            if 0 <= i && i < n {
                h = append(h, a[i]) //build the heap
                if i == n-1 {
                    for j := n - 1; j >= 0; j-- {
                        HeapAdjust(h, j, n)
                        fmt.Println(h)
                    }
                }
            } else {
                //算法开始
                if a[i].Exp > h[0].Exp {
                    h[0].Exp = a[i].Exp
                    p := 0 //下标
                    HeapAdjust(h, p, n)
                }
                //算法结束
            }
        }
    
        return h, nil
    }
    
    //h是待调整的堆数组，p是待调整的数组元素的位置，n是数组的长度
    //本函数功能是：根据数组h构建min根堆
    func HeapAdjust(h []leveldb.UserExp, p, n int) {
        //下标
        for p < n { //排序条件 不能超出n
            q := p<<1 | 1 //h[p]的左子树
            if q >= n {
                break //已经大于最大的树的节点的下标[0,n-1]
            }
            if (q < n-1) && (h[q+1].Exp < h[q].Exp) { //q<n-1 确保q+1不超出不超出n-1 左大于右
                q = q + 1 //用相邻的两个值的小值的下标
            }
    
            if h[q].Exp < h[p].Exp { //用最小的值比较 确保最小的值放在数组的最前面
                t := h[p] //交换
                h[p] = h[q]
                h[q] = t
                p = q //存的是最小的值
            } else {
                break
            }
        }
    }