Use Priority queue

Signed-off-by: Leo Ma <begeekmyfriend@192.168.1.5>

Author: Leo Ma

023_merge_k_sorted_lists/merge_lists.c

@@ -1,45 +1,144 @@
 #include <stdio.h>
 #include <stdlib.h>
-#include <limits.h>
+#include <string.h>
 
 struct ListNode {
     int val;
     struct ListNode *next;
 };
 
+struct PriorityQueue {
+    struct ListNode **nodes;
+    int size;
+};
+
+static inline void swap(struct ListNode **a, struct ListNode **b)
+{
+    struct ListNode *tmp = *a;
+    *a = *b;
+    *b = tmp;
+}
+
+static inline int left(int i) { return i * 2 + 1; }
+static inline int right(int i) { return left(i) + 1; }
+static inline int parent(int i) { return (i - 1) / 2; }
+
+static void queue_dump(struct PriorityQueue *queue)
+{
+    int i;
+    for (i = 0; i < queue->size; i++) {
+        printf("%d ", queue->nodes[i]->val);
+    }
+    printf("\n");
+}
+
+static void percolate_up(struct ListNode **nodes, int i)
+{
+    while (i >= 0 && nodes[parent(i)]->val > nodes[i]->val) {
+        swap(nodes + parent(i), nodes + i);
+        i = parent(i);
+    }
+}
+
+static void percolate_down1(struct ListNode **nodes, int size, int child)
+{
+    int i, min;
+    for (i = child; i >= 0; i = parent(i)) {
+        if (right(i) < size) {
+            min = nodes[left(i)]->val < nodes[right(i)]->val ? left(i) : right(i);
+        } else {
+            min = left(i);
+        }
+        if (nodes[min]->val < nodes[i]->val) {
+            swap(nodes + min, nodes + i);
+        } else {
+            break;
+        }
+    }
+}
+
+static void percolate_down2(struct ListNode **nodes, int size)
+{
+    int i, min;
+    for (i = 0; left(i) < size; i = min) {
+        if (right(i) < size) {
+            min = nodes[left(i)]->val < nodes[right(i)]->val ? left(i) : right(i);
+        } else {
+            min = left(i);
+        }
+        if (nodes[min]->val < nodes[i]->val) {
+            swap(nodes + min, nodes + i);
+        } else {
+            break;
+        }
+    }
+}
+
+static void heap_build(struct PriorityQueue *queue)
+{
+    int i;
+    for (i = queue->size / 2 - 1; i > 0; i--) {
+        percolate_down1(queue->nodes, queue->size, i);
+    }
+}
+
+static void put(struct PriorityQueue *queue, struct ListNode *node)
+{
+    queue->nodes[queue->size++] = node;
+    percolate_up(queue->nodes, queue->size - 1);
+}
+
+static struct ListNode *get(struct PriorityQueue *queue)
+{
+    int i;
+    struct ListNode *p = queue->nodes[0];
+    swap(queue->nodes, queue->nodes + queue->size - 1);
+    queue->size--;
+    percolate_down2(queue->nodes, queue->size);
+    return p;
+}
+
+static struct PriorityQueue *init(int size)
+{
+    struct PriorityQueue *queue = malloc(sizeof(*queue));
+    queue->nodes = malloc(size * sizeof(*queue->nodes));
+    queue->size = 0;
+    return queue;
+}
+
 static struct ListNode* mergeKLists(struct ListNode** lists, int listsSize)
 {
     if (listsSize == 0) {
         return NULL;
     }
+
     if (listsSize == 1) {
         return lists[0];
     }
 
-    int i, index;
-    struct ListNode dummy, *p, *prev;
+    int i;
+    struct ListNode dummy;
+    struct ListNode *prev;
+    struct PriorityQueue *queue = init(listsSize);
 
     dummy.next = NULL;
     prev = &dummy;
-    index = 0;
-    while (index != -1) {
-        int min = INT_MAX;
-        index = -1;
-        for (i = 0; i < listsSize; i++) {
-            if (lists[i] != NULL && lists[i]->val < min) {
-                min = lists[i]->val;
-                index = i;
-            }
+
+    for (i = 0; i < listsSize; i++) {
+        if (lists[i] != NULL) {
+            put(queue, lists[i]);
         }
+    }
+    heap_build(queue);
 
-        if (index != -1) {
-            p = malloc(sizeof(*p));
-            p->val = min;
-            p->next = NULL;
-            prev->next = p;
-            prev = p;
-            lists[index] = lists[index]->next;
+    while (queue->size > 0) {
+        struct ListNode *n = get(queue);
+        prev->next = n;
+        prev = n;
+        if (n->next != NULL) {
+            put(queue, n->next);
         }
+        n->next = NULL;
     }
 
     return dummy.next;