Constant extra space merge

gistfile1.rs

// https://twitter.com/mraleph/status/559477604659773440
// merge two sorted arrays in O(1) memory and O(n) time. nice excercise
// sort array of length k + m such that a[0] <= ... <= a[k-1] & a[k] <= ... <= a[k+m-1], in place in O(k+m) time

#![allow(unstable)]

// data.len() == k + m
fn merge(data: &mut [u8], k: usize) {
    // M = data[..a] is merged
    // A = data[a..b] is initially the first sub-array, of length k
    // B = data[b..c] is initially empty
    // C = data[c..] is initially the second sub-array of length m
    let mut a = 0;
    let mut b = k;
    let mut c = k;

    loop {
//        println!("{:?} {:?} {:?} {:?}", &data[..a], &data[a..b], &data[b..c], &data[c..]);
        assert!(data[..a].is_sorted());
        assert!(data[a..b].is_sorted());
        assert!(data[b..c].is_sorted());
        assert!(data[c..].is_sorted());

        let a_non_empty = a < b;
        let b_non_empty = b < c;
        let c_non_empty = c < data.len();
        if a_non_empty && b_non_empty && c_non_empty {
            // Three-way merge
            if data[a] > data[b] {
                if data[b] > data[c] {
                    data.swap(a, c);
                    c += 1;
                } else {
                    data.swap(a, b);
                    // :(
                    for i in b..(c - 1) {
                        data.swap(i, i + 1);
                    }
                }
            } else if data[a] > data[c] {
                data.swap(a, c);
                c += 1;
            }
            a += 1;
        } else if a_non_empty && b_non_empty {
            // Two-way merge
            if data[a] > data[b] {
                data.swap(a, b);
                // Truncate B to 1
                // and make the empty C be the rest of B.
                c = b + 1;
            }
            a += 1;
        } else if a_non_empty && c_non_empty {
            // Two-way merge
            if data[a] > data[c] {
                // Insert data[a] into the empty B.
                data.swap(a, c);
                c += 1;
            }
            a += 1;
        } else if b_non_empty && c_non_empty {
            assert!(a == b);
            // Two-way merge
            if data[a] > data[c] {
                // Make the empty A be B.
                b = c;
                // Insert data[a] into the now-empty B.
                data.swap(a, b);
                c += 1;
            } else {
                b += 1;
            }
            a += 1;
        } else {
            // Only one (or zero) of the three sub-arrays remaining
            // to merge with itself: we’re done.
            break
        }
    }
}

#[test]
fn it_works() {
    use std::rand::{Rng, StdRng};
    let mut rng = StdRng::new().unwrap();
    for k in 0us..20 {
        for m in 0us..20 {
            for _ in 0..100 {
                let mut data = rng.gen_iter().take(k + m).collect::<Vec<_>>();

                let mut sorted = data.clone();
                sorted.sort();

                data[..k].sort();
                data[k..].sort();
                merge(&mut *data, k);

                assert!(data.is_sorted());
            }
        }
    }
}


trait IsSorted {
    fn is_sorted(&mut self) -> bool;
}

impl<T: Ord> IsSorted for [T] {
    fn is_sorted(&mut self) -> bool {
        self.windows(2).all(|chunk| match chunk {
            [ref a, ref b] => a <= b,
            _ => panic!("Unexpected window size")
        })
    }
}