Sort Colors

Problem statement

nums only contains 0, 1, and 2. Sort the array in place into that same order (all zeros, then ones, then twos).

Example:

Input:

[2, 0, 2, 1, 1, 0]

Expected output:

[0, 0, 1, 1, 2, 2]

Practice on LeetCode: Sort Colors

Golang Solution

func sortColors(nums []int) {
    if len(nums) <= 1 {
        return
    }
    i, j := 0, 0
    k := len(nums) - 1
    for j <= k {
        if nums[j] == 0 {
            nums[i], nums[j] = nums[j], nums[i]
            i += 1
            j += 1
        } else if nums[j] == 2 {
            nums[k], nums[j] = nums[j], nums[k]
            k -= 1
        } else {
            j += 1
        }
    }
}