From ea1d62bd923296c61ee17da2ecef9ed6f2f8c0d0 Mon Sep 17 00:00:00 2001
From: ameanasad <amean.asad1999@gmail.com>
Date: Mon, 18 Sep 2023 17:22:52 -0400
Subject: [PATCH] fix: top N selection

---
 node_heap.go | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/node_heap.go b/node_heap.go
index 6067abf..1d6092f 100644
--- a/node_heap.go
+++ b/node_heap.go
@@ -51,13 +51,22 @@ func (nh *NodeHeap) PeekRandom() *Node {
 
 func (nh *NodeHeap) TopN(n int) []*Node {
 	m := make([]*Node, 0, n)
+	temp := make([]*Node, 0, n)
 	nh.lk.Lock()
 	defer nh.lk.Unlock()
+
 	heap.Init(nh)
-	for i := 0; i < n && i < len(nh.Nodes); i++ {
-		node := nh.Nodes[i]
+	for i := 0; i < n && nh.Len() > 0; i++ {
+		item := heap.Pop(nh)
+		node := item.(*Node)
 		m = append(m, node)
+		temp = append(temp, node)
+	}
+
+	for _, node := range temp {
+		heap.Push(nh, node)
 	}
+
 	return m
 }