Discussion on: Merge k sorted arrays in Rust

[Caleb Sander]

This is a good problem, got asked it in a Facebook interview! A few suggestions on the implementation:

• It will panic if any of the input arrays are empty
• get_item() does a bounds check every time it's called, which shouldn't be necessary. You can store the next element in the Item struct to fix this.
• You can avoid checking heap.is_empty() and using .unwrap() if you use while let Some(it) = heap.pop() instead
• You can use heap.peek_mut() to avoid sifting the heap twice each time through the loop
• You can pre-allocate sorted with the correct capacity

Here's what I came up with:

struct MergeCandidate<'a, T> {
    first: T,
    rest: &'a [T],
}

impl <'a, T: Copy> From<(&T, &'a [T])> for MergeCandidate<'a, T> {
    fn from((&first, rest): (&T, &'a [T])) -> Self {
        MergeCandidate { first, rest }
    }
}

impl <T: PartialEq> PartialEq for MergeCandidate<'_, T> {
    fn eq(&self, other: &Self) -> bool {
        self.first == other.first
    }
}

impl <T: Eq> Eq for MergeCandidate<'_, T> {}

impl <T: PartialOrd> PartialOrd for MergeCandidate<'_, T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.first.partial_cmp(&self.first)
    }
}

impl <T: Ord> Ord for MergeCandidate<'_, T> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.first.cmp(&self.first)
    }
}

pub fn merge<T: Ord + Copy>(lists: &[&[T]]) -> Vec<T> {
    let mut candidates = BinaryHeap::with_capacity(lists.len());
    let mut total_length = 0;
    for &list in lists {
        if let Some(candidate) = list.split_first() {
            candidates.push(MergeCandidate::from(candidate));
            total_length += list.len();
        }
    }
    let mut sorted = Vec::with_capacity(total_length);
    while let Some(mut merge_candidate) = candidates.peek_mut() {
        sorted.push(merge_candidate.first);
        if let Some(candidate) = merge_candidate.rest.split_first() {
            *merge_candidate = MergeCandidate::from(candidate);
        }
        else {
            PeekMut::pop(merge_candidate);
        }
    }
    sorted
}

[creativcoder]

Hey Caleb,

Thank you for improving the solution. The idea to store the rest of items in Item is neat!