LeetCode-in-Scala.github.io

75. Sort Colors

Medium

Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.

We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.

You must solve this problem without using the library’s sort function.

Example 1:

Input: nums = [2,0,2,1,1,0]

Output: [0,0,1,1,2,2]

Example 2:

Input: nums = [2,0,1]

Output: [0,1,2]

Example 3:

Input: nums = [0]

Output: [0]

Example 4:

Input: nums = [1]

Output: [1]

Constraints:

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

Solution

object Solution {
    def sortColors(nums: Array[Int]): Unit = {
        var zeroes = 0
        var ones = 0

        for (i <- nums.indices) {
            if (nums(i) == 0) {
                nums(zeroes) = 0
                zeroes += 1
            } else if (nums(i) == 1) {
                ones += 1
            }
        }

        for (j <- zeroes until zeroes + ones) {
            nums(j) = 1
        }

        for (k <- zeroes + ones until nums.length) {
            nums(k) = 2
        }
    }
}