Rust Gym Leetcode 23 Merge k Sorted Lists Rust Solution

23. Merge k Sorted Lists

You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

Example 1:

Input: lists = [[1,4,5],[1,3,4],[2,6]]
Output: [1,1,2,3,4,4,5,6]
Explanation: The linked-lists are:
[
  1->4->5,
  1->3->4,
  2->6
]
merging them into one sorted list:
1->1->2->3->4->4->5->6

Example 2:

Input: lists = []
Output: []

Example 3:

Input: lists = [[]]
Output: []

Constraints:

k == lists.length
0 <= k <= 10^4
0 <= lists[i].length <= 500
-10^4 <= lists[i][j] <= 10^4
lists[i] is sorted in ascending order.
The sum of lists[i].length won't exceed 10^4.

23. Merge k Sorted Lists

struct Solution;
use rustgym_util::*;
use std::collections::VecDeque;

impl Solution {
    fn merge_k_lists(lists: Vec<ListLink>) -> ListLink {
        if lists.is_empty() {
            return None;
        }
        let mut queue: VecDeque<ListLink> = lists.into_iter().collect();
        while queue.len() > 1 {
            let merged_list = Self::merge(queue.pop_front().unwrap(), queue.pop_front().unwrap());
            queue.push_back(merged_list);
        }
        queue.pop_back().unwrap()
    }

    fn merge(a: ListLink, b: ListLink) -> ListLink {
        if a.is_none() && b.is_none() {
            return None;
        }
        if a.is_none() {
            return b;
        }
        if b.is_none() {
            return a;
        }
        let mut a = a.unwrap();
        let mut b = b.unwrap();
        if a.val < b.val {
            a.next = Self::merge(a.next.take(), Some(b));
            Some(a)
        } else {
            b.next = Self::merge(Some(a), b.next.take());
            Some(b)
        }
    }
}

#[test]
fn test() {
    let lists = vec![list!(1, 4, 5), list!(1, 3, 4), list!(2, 6)];
    let res = list!(1, 1, 2, 3, 4, 4, 5, 6);
    assert_eq!(Solution::merge_k_lists(lists), res);
}