feat: refactor top_k.ts (krahets#899)

Gonglja · Oct 28, 2023 · c37f098 · c37f098
1 parent f98b337
commit c37f098
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 40 deletions.
diff --git a/codes/typescript/chapter_heap/my_heap.ts b/codes/typescript/chapter_heap/my_heap.ts
@@ -122,32 +122,34 @@ class MaxHeap {
 }
 
 /* Driver Code */
-/* 初始化大顶堆 */
-const maxHeap = new MaxHeap([9, 8, 6, 6, 7, 5, 2, 1, 4, 3, 6, 2]);
-console.log('\n输入列表并建堆后');
-maxHeap.print();
-
-/* 获取堆顶元素 */
-let peek = maxHeap.peek();
-console.log(`\n堆顶元素为 ${peek}`);
-
-/* 元素入堆 */
-const val = 7;
-maxHeap.push(val);
-console.log(`\n元素 ${val} 入堆后`);
-maxHeap.print();
-
-/* 堆顶元素出堆 */
-peek = maxHeap.pop();
-console.log(`\n堆顶元素 ${peek} 出堆后`);
-maxHeap.print();
-
-/* 获取堆大小 */
-const size = maxHeap.size();
-console.log(`\n堆元素数量为 ${size}`);
-
-/* 判断堆是否为空 */
-const isEmpty = maxHeap.isEmpty();
-console.log(`\n堆是否为空 ${isEmpty}`);
+if (require.main === module) {
+    /* 初始化大顶堆 */
+    const maxHeap = new MaxHeap([9, 8, 6, 6, 7, 5, 2, 1, 4, 3, 6, 2]);
+    console.log('\n输入列表并建堆后');
+    maxHeap.print();
+
+    /* 获取堆顶元素 */
+    let peek = maxHeap.peek();
+    console.log(`\n堆顶元素为 ${peek}`);
+
+    /* 元素入堆 */
+    const val = 7;
+    maxHeap.push(val);
+    console.log(`\n元素 ${val} 入堆后`);
+    maxHeap.print();
+
+    /* 堆顶元素出堆 */
+    peek = maxHeap.pop();
+    console.log(`\n堆顶元素 ${peek} 出堆后`);
+    maxHeap.print();
+
+    /* 获取堆大小 */
+    const size = maxHeap.size();
+    console.log(`\n堆元素数量为 ${size}`);
+
+    /* 判断堆是否为空 */
+    const isEmpty = maxHeap.isEmpty();
+    console.log(`\n堆是否为空 ${isEmpty}`);
+}
 
 export { MaxHeap };
diff --git a/codes/typescript/chapter_heap/top_k.ts b/codes/typescript/chapter_heap/top_k.ts
@@ -6,25 +6,49 @@
 
 import { MaxHeap } from './my_heap';
 
+/* 元素入堆 */
+function pushMinHeap(maxHeap: MaxHeap, val: number): void {
+    // 元素取反
+    maxHeap.push(-val);
+}
+
+/* 元素出堆 */
+function popMinHeap(maxHeap: MaxHeap): number {
+    // 元素取反
+    return -maxHeap.pop();
+}
+
+/* 访问堆顶元素 */
+function peekMinHeap(maxHeap: MaxHeap): number {
+    // 元素取反
+    return -maxHeap.peek();
+}
+
+/* 取出堆中元素 */
+function getMinHeap(maxHeap: MaxHeap): number[] {
+    // 元素取反
+    return maxHeap.getMaxHeap().map((num: number) => -num);
+}
+
 /* 基于堆查找数组中最大的 k 个元素 */
 function topKHeap(nums: number[], k: number): number[] {
-    // 将堆中所有元素取反，从而用大顶堆来模拟小顶堆
-    const invertedNums = nums.map((num) => -num);
+    // 初始化小顶堆
+    // 请注意：我们将堆中所有元素取反，从而用大顶堆来模拟小顶堆
+    const maxHeap = new MaxHeap([]);
     // 将数组的前 k 个元素入堆
-    const heap = new MaxHeap(invertedNums.slice(0, k));
+    for (let i = 0; i < k; i++) {
+        pushMinHeap(maxHeap, nums[i]);
+    }
     // 从第 k+1 个元素开始，保持堆的长度为 k
-    for (let i = k; i < invertedNums.length; i++) {
-        // 若当前元素小于堆顶元素，则将堆顶元素出堆、当前元素入堆
-        if (invertedNums[i] < heap.peek()) {
-            heap.pop();
-            heap.push(invertedNums[i]);
+    for (let i = k; i < nums.length; i++) {
+        // 若当前元素大于堆顶元素，则将堆顶元素出堆、当前元素入堆
+        if (nums[i] > peekMinHeap(maxHeap)) {
+            popMinHeap(maxHeap);
+            pushMinHeap(maxHeap, nums[i]);
         }
     }
-    // 取出堆中元素
-    const maxHeap = heap.getMaxHeap();
-    // 对堆中元素取相反数
-    const invertedMaxHeap = maxHeap.map((num) => -num);
-    return invertedMaxHeap;
+    // 返回堆中元素
+    return getMinHeap(maxHeap);
 }
 
 /* Driver Code */