leetcode/
merge_k_sorted_lists.rs

1//! # 23. 合并K个升序链表
2//!
3//! 难度 困难
4//!
5//! 给你一个链表数组，每个链表都已经按升序排列。
6//!
7//! 请你将所有链表合并到一个升序链表中，返回合并后的链表。
8//!
9//! ## 示例 1：
10//!
11//! ```text
12//! 输入：lists = [[1,4,5],[1,3,4],[2,6]]
13//! 输出：[1,1,2,3,4,4,5,6]
14//! 解释：链表数组如下：
15//! [
16//!   1->4->5,
17//!   1->3->4,
18//!   2->6
19//! ]
20//! 将它们合并到一个有序链表中得到。
21//! 1->1->2->3->4->4->5->6
22//! ```
23//!
24//! ## 示例 2：
25//!
26//! ```text
27//! 输入：lists = []
28//! 输出：[]
29//! ```
30//!
31//! ## 示例 3：
32//!
33//! ```text
34//! 输入：lists = [[]]
35//! 输出：[]
36//! ```
37//!
38//! ## 提示：
39//!
40//! - `k` == `lists.length`
41//! - `0` <= `k` <= `10^4`
42//! - `0` <= `lists[i].length` <= `500`
43//! - `-10^4` <= `lists[i][j]` <= `10^4`
44//! - `lists[i]` 按 **升序** 排列
45//! - `lists[i].length` 的总和不超过 `10^4`
46//!
47//! See [leetcode](https://leetcode-cn.com/problems/merge-k-sorted-lists/)
48
49// use crate::ListNode;
50// Definition for singly-linked list.
51#[derive(PartialEq, Eq, Clone, Debug)]
52pub struct ListNode {
53  pub val: i32,
54  pub next: Option<Box<ListNode>>
55}
56
57impl ListNode {
58  #[inline]
59  pub fn new(val: i32) -> Self {
60    ListNode {
61      next: None,
62      val
63    }
64  }
65}
66
67use std::cmp::{Ord, Ordering, PartialEq, Reverse};
68use std::collections::BinaryHeap;
69
70impl Ord for ListNode {
71    fn cmp(&self, other: &Self) -> Ordering {
72        self.val.cmp(&other.val)
73    }
74}
75
76impl PartialOrd for ListNode {
77    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
78        Some(self.cmp(other))
79    }
80}
81
82pub struct Solution;
83
84impl Solution {
85
86    /// 使用 BinaryHeap，但需要转化为小顶堆
87    pub fn merge_k_lists(lists: Vec<Option<Box<ListNode>>>) -> Option<Box<ListNode>> {
88        if lists.is_empty() {
89            return None;
90        }
91
92        let mut priority_queue = BinaryHeap::with_capacity(lists.len());
93
94        for node in lists {
95            if let Some(node) = node {
96                priority_queue.push(Reverse(node))
97            }
98        }
99
100        let mut result = Box::new(ListNode::new(0));
101        let mut ptr = &mut result;
102        while let Some(Reverse(mut node)) = priority_queue.pop() {
103            if let Some(next) = node.next.take() {
104                priority_queue.push(Reverse(next));
105            }
106            ptr.next = Some(node);
107            ptr = ptr.next.as_mut().unwrap()
108        }
109
110        result.next
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::{Solution, ListNode};
117
118    fn slice_to_list(a: &[i32]) -> Option<Box<ListNode>> {
119        a.iter().fold(None, |mut list, &v| {
120            let node = Box::new(ListNode::new(v));
121            if let Some(ref mut a) = list {
122                let mut last = &mut a.next;
123                while let Some(ref mut node) = last {
124                    last = &mut node.next;
125                }
126                *last = Some(node);
127                list
128            } else {
129                Some(node)
130            }
131        })
132    }
133
134    fn list_to_slice(mut a: Option<Box<ListNode>>) -> Vec<i32> {
135        std::iter::from_fn(move || {
136            if let Some(mut v) = a.take() {
137                let item = Some(v.val);
138                a = v.next.take();
139                item
140            } else {
141                None
142            }
143        }).collect()
144    }
145
146    #[test]
147    fn test() {
148        let lists: Vec<_> = vec![vec![1,4,5], vec![1,3,4], vec![2,6]].into_iter().map(|a| slice_to_list(&a[..])).collect();
149        assert_eq!(&[1i32,1,2,3,4,4,5,6][..], list_to_slice(Solution::merge_k_lists(lists)));
150
151        let lists: Vec<_> = vec![].into_iter().map(|a: Vec<i32>| slice_to_list(&a[..])).collect();
152        let expect: &[i32] = &[][..];
153        assert_eq!(expect, list_to_slice(Solution::merge_k_lists(lists)));
154
155        let lists: Vec<_> = vec![vec![]].into_iter().map(|a| slice_to_list(&a[..])).collect();
156        let expect: &[i32] = &[][..];
157        assert_eq!(expect, list_to_slice(Solution::merge_k_lists(lists)));
158    }
159}
leetcode/merge_k_sorted_lists.rs

leetcode/
merge_k_sorted_lists.rs
