Sort a LL of 0s, 1s and 2s

The problem, simply

You get a linked list where every box holds a 0, a 1 or a 2. Sort it so all the 0s come first, then the 1s, then the 2s. Give back the new head.

There are only three kinds of numbers, so you never need a real comparison sort. The question is whether you copy numbers or move boxes.

A story to remember it

Picture it

Examples

Way 1: the simple way

Count, then overwrite.

1. Walk the chain and count how many 0s, 1s and 2s there are.
2. Walk again from the head. Write 0 into the first count0 boxes, then 1 into the next count1 boxes, then 2 into the rest.

class Solution {
    public Node segregate(Node head) {
        int[] count = new int[3];
        for (Node cur = head; cur != null; cur = cur.next) count[cur.data]++;
        Node cur = head;
        for (int d = 0; d < 3; d++) {
            for (int i = 0; i < count[d]; i++) {   // write d, count[d] times
                cur.data = d;
                cur = cur.next;
            }
        }
        return head;
    }
}

How long it takes: two walks, O(n). Extra room: three counters, O(1). This is often the expected answer for this exact problem. But if each box carries more than a number (a name, a record), overwriting the number would mix up the records, so interviewers then ask for the next way.

Way 2: the clever way

Move the real boxes into three lines, then join them.

1. Make three dummy boxes, one per line, and a tail finger for each: zero, one, two.
2. Walk the chain. Hook each box to the end of its line and move that tail.
3. Join: the end of the zeros points to the first 1, or to the first 2 if there are no 1s. The end of the ones points to the first 2.
4. Set the end of the twos to null.
5. Give back the box after the zeros dummy.

class Solution {
    public Node segregate(Node head) {
        Node zeroHead = new Node(0), oneHead = new Node(0), twoHead = new Node(0);
        Node zero = zeroHead, one = oneHead, two = twoHead;    // tails of the three lines
        for (Node cur = head; cur != null; cur = cur.next) {
            if (cur.data == 0) { zero.next = cur; zero = cur; }
            else if (cur.data == 1) { one.next = cur; one = cur; }
            else { two.next = cur; two = cur; }
        }
        zero.next = (oneHead.next != null) ? oneHead.next : twoHead.next;  // skip an empty ones line
        one.next = twoHead.next;
        two.next = null;                                       // cut the old arrow
        return zeroHead.next;                                  // works even with no 0s
    }
}

How long it takes: one walk, O(n). Extra room: three dummy boxes and three tails, O(1). No number is overwritten, and boxes with the same number keep their old order.

Watch out for

• Forgetting two.next = null. The last 2 still points to its old neighbor, and the chain loops.
• Returning head. The old head may be a 2 that is now near the end. Return zeroHead.next.
• No 0s at all: zero is still the dummy, so zeroHead.next becomes the ones line (or the twos line). The code needs no special case.

What interviewers ask next

• Do it on an array (the Dutch national flag problem). Use three fingers low, mid, high and swap in place: a 0 swaps to low, a 2 swaps to high, a 1 just moves mid on.
• Why not use the array trick on a list? It needs to swap with the back end, and a singly list cannot walk backwards. Three lines are the natural list version.
• What if values go from 0 to m? Use m + 1 lines (or m + 1 counters). That is counting sort, O(n + m).
• Make it one walk with counting? Counting needs two walks, because you must know the counts before writing. The three lines way is a single walk.

Remember this

Grown-up words: the bins are three sublists with dummy head nodes, counting and overwriting is counting sort, and the array version with three fingers is the Dutch national flag algorithm.