revise

heap_sort.py

@@ -0,0 +1,49 @@
+"""
+在海量数据中选取topK个数据:
+整个操作中,遍历数组需要O(n)的时间复杂度,每次调整小顶堆的时间复杂度是O(logK),加起来就是O(nlogK)的复杂度;
+如果K远小于n的话,O(nlogK)其实就接近于O(n),甚至会更快,因此也是十分高效的.
+"""
+
+# 构建大顶堆
+def heapify(arr, n, i, dim):
+    largest = i
+    left = 2*i + 1
+    right = 2*i + 2
+    # 与左节点进行比较
+    if left < n and arr[i][dim] < arr[left][dim]:
+        largest = left
+    # 与左节点比较后再与右节点进行比较
+    if right < n and arr[largest][dim] < arr[right][dim]:
+        largest = right
+    # 通过上面跟左右节点比较后，得出三个元素之间较大的下标
+    # 若较大下标不是父节点的下标，说明交换后需要重新调整大顶堆
+    if largest != i:
+        arr[i], arr[largest] = arr[largest], arr[i]
+        # 重新调整大顶堆
+        heapify(arr, n, largest, dim)
+
+# 堆排序
+def heap_sort(arr, topk=None, dim=1): 
+    n = len(arr) 
+    # 构造大顶堆，从非叶子节点开始倒序遍历，因此是l//2 -1 就是最后一个非叶子节点
+    for i in range(n//2-1, -1, -1):
+        heapify(arr, n, i, dim)
+    # 若topk不为None,则进行设定
+    topk = n if topk > n or topk is None else topk
+    res_list = []
+    # 上面的循环完成了大顶堆的构造，那么就开始把根节点跟末尾节点交换，然后重新调整大顶堆
+    for i in range(n-1, n-1-topk, -1):
+        res_list.append(arr[0])
+        arr[0] = arr[i]
+        # arr[i], arr[0] = arr[0], arr[i]
+        heapify(arr, i, 0, dim)
+    
+    return res_list
+    # return arr[n-topk:][::-1]
+
+if __name__ == "__main__":
+    arr = [[0,12], [1,11], [2,13], [3,5], [4,6], [5,7]]
+    arr = heap_sort(arr, 3)
+    print ("排序后")
+    for i in range(len(arr)):
+        print(arr[i])

info_retrieval.py

@@ -8,6 +8,7 @@ import sqlite3
 from config import sqlite_path, sqlite_copy_path, log_msg
 from data_preprocessing import DataPreProcessing
 from word_segment import Word_Segment
+from heap_sort import heap_sort
 
 class Info_Retrieval():
     def __init__(self, data_process, logger=None, n_grams_flag=False):