modify code

algorithmzuo · algorithmzuo · commit d773fe8ca8f4 · 2021-01-05T17:20:56.000+08:00
diff --git a/src/class05/Code06_MergeKSortedLists.java b/src/class05/Code06_MergeKSortedLists.java
@@ -0,0 +1,51 @@
+package class05;
+
+import java.util.Comparator;
+import java.util.PriorityQueue;
+
+public class Code06_MergeKSortedLists {
+
+	public static class ListNode {
+		public int val;
+		public ListNode next;
+	}
+
+	public static class ListNodeComparator implements Comparator<ListNode> {
+
+		@Override
+		public int compare(ListNode o1, ListNode o2) {
+			return o1.val - o2.val;
+		}
+
+	}
+
+	public static ListNode mergeKLists(ListNode[] lists) {
+		if (lists == null) {
+			return null;
+		}
+		PriorityQueue<ListNode> heap = new PriorityQueue<>(new ListNodeComparator());
+		for (int i = 0; i < lists.length; i++) {
+			if (lists[i] != null) {
+				heap.add(lists[i]);
+			}
+		}
+		if (heap.isEmpty()) {
+			return null;
+		}
+		ListNode head = heap.poll();
+		ListNode pre = head;
+		if (pre.next != null) {
+			heap.add(pre.next);
+		}
+		while (!heap.isEmpty()) {
+			ListNode cur = heap.poll();
+			pre.next = cur;
+			pre = cur;
+			if (cur.next != null) {
+				heap.add(cur.next);
+			}
+		}
+		return head;
+	}
+
+}
