Topics › Linked List › Medium Problems of LL

Sort a Linked List of 0's 1's and 2's

Medium TUF ID 629 A2Z

Problem

You are given the head of a singly linked list in which every node holds one of the values 0, 1 or 2. Rearrange the list so that all 0s come first, then all 1s, then all 2s, and return the head. A general comparison sort would take O(n log n); with only three distinct values the job can be done in O(n) time, and the preferred solution does it by re-linking nodes in a single pass without changing any value.

Constraints: 0 ≤ n ≤ 10^5 nodes, every value is 0, 1 or 2.

static class ListNode {
    int val;
    ListNode next;
    ListNode(int val) { this.val = val; }
    ListNode(int val, ListNode next) { this.val = val; this.next = next; }
}

Example 1

head = 1 -> 0 -> 2 -> 1 -> 0 -> 2 -> 1
answer = 0 -> 0 -> 1 -> 1 -> 1 -> 2 -> 2

Example 2

head = 2 -> 2 -> 0
answer = 0 -> 2 -> 2            // there are no 1s at all

Example 3

head = 1
answer = 1

Approach 1: Brute force, count and overwrite

Intuition. Since there are only three possible values, the sorted list is completely described by three numbers: how many 0s, how many 1s, how many 2s. Count them in one pass, then walk the list again and write 0 into the first zeros nodes, 1 into the next ones nodes, and 2 into the rest. This is counting sort with three buckets.

Algorithm.

  1. Walk the list and count zeros, ones, twos.
  2. Walk the list again. At each node: if zeros > 0, write 0 and decrement zeros; else if ones > 0, write 1 and decrement ones; else write 2.
  3. Return head.
static ListNode bruteForce(ListNode head) {
    int zeros = 0, ones = 0, twos = 0;
    for (ListNode t = head; t != null; t = t.next) {
        if (t.val == 0) zeros++;
        else if (t.val == 1) ones++;
        else twos++;
    }
    for (ListNode t = head; t != null; t = t.next) {
        if (zeros > 0) { t.val = 0; zeros--; }
        else if (ones > 0) { t.val = 1; ones--; }
        else t.val = 2;
    }
    return head;
}

Complexity. Time O(n): two passes. Space O(1): three counters.

This is already linear and constant-space, so why look further? Two reasons. It takes two passes, and it changes the values stored in nodes rather than moving the nodes. In many real settings the node is an object with other fields, and overwriting its payload is not an option. The interview expectation is the single-pass re-linking solution.

Approach 2: Optimal, three chains joined together

Intuition. Keep three separate chains while walking the list once: one for nodes with value 0, one for 1, one for 2. Each node is appended to the chain matching its value. At the end, link the tail of the 0 chain to the head of the 1 chain, the tail of the 1 chain to the head of the 2 chain, and terminate the 2 chain with null. Dummy head nodes for each chain remove all the special cases around “is this chain still empty?”.

The only subtlety is an empty middle chain: if there are no 1s, the 0 chain must link directly to the 2 chain. The dummy node makes this a one-line check.

Algorithm.

  1. Create dummy nodes zeroDummy, oneDummy, twoDummy and tail pointers zero, one, two initialised to them.
  2. For each node t in the list: append t to the chain matching t.val and advance that chain’s tail.
  3. Join: zero.next = (oneDummy.next != null) ? oneDummy.next : twoDummy.next; one.next = twoDummy.next; two.next = null.
  4. Return zeroDummy.next.
static ListNode optimal(ListNode head) {
    ListNode zeroDummy = new ListNode(-1), oneDummy = new ListNode(-1), twoDummy = new ListNode(-1);
    ListNode zero = zeroDummy, one = oneDummy, two = twoDummy;
    for (ListNode t = head; t != null; t = t.next) {
        if (t.val == 0) { zero.next = t; zero = t; }
        else if (t.val == 1) { one.next = t; one = t; }
        else { two.next = t; two = t; }
    }
    zero.next = (oneDummy.next != null) ? oneDummy.next : twoDummy.next;   // skip the 1-chain if it is empty
    one.next = twoDummy.next;
    two.next = null;                                                        // the last 2 (or the last node overall) ends the list
    return zeroDummy.next;
}

Dry run on 1 -> 0 -> 2 -> 1 -> 0 -> 2 -> 1:

node 0 chain 1 chain 2 chain
1   1  
0 0 1  
2 0 1 2
1 0 1 -> 1 2
0 0 -> 0 1 -> 1 2
2 0 -> 0 1 -> 1 2 -> 2
1 0 -> 0 1 -> 1 -> 1 2 -> 2
join 0 -> 0 -> 1 -> 1 -> 1 -> 2 -> 2 -> null    

Complexity. Time O(n): one pass plus three constant-time joins. Space O(1): three dummy nodes and three pointers.

The Java tests compare both methods against Arrays.sort on the same inputs and, for the optimal method, verify that the result is made of exactly the original node objects (same count and every original node present). Inputs include lists with no 0s, no 1s, and no 2s, since each missing group exercises a different join.

Edge cases and pitfalls