LeetCode第75题 - 颜色分类

题目

解答
两趟扫描输入数组。

package leetcode.editor.cn;

//leetcode submit region begin(Prohibit modification and deletion)
class Solution {
    public void sortColors(int[] nums) {
        if (nums == null || nums.length == 1) {
            return;
        }

        int red = 0;
        int white = 0;
        int blue = 0;
        for (int i = 0; i < nums.length; ++i) {
            switch (nums[i]) {
                case 0:
                    ++red;
                    break;
                case 1:
                    ++white;
                    break;
                case 2:
                    ++blue;
                    break;
                default:
            }
        }

        for (int i = 0; i < nums.length; ++i) {
            if (red != 0) {
                nums[i] = 0;
                --red;
            } else if (white != 0) {
                nums[i] = 1;
                --white;
            } else if (blue != 0) {
                nums[i] = 2;
                --blue;
            }
        }
    }

}
//leetcode submit region end(Prohibit modification and deletion)

测试用例

package leetcode.editor.cn;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;

public class SolutionTest {
    private Solution s;

    @Before
    public void setUp() throws Exception {
        s = new Solution();
    }

    @Test
    public void test1() {
        int[] input = new int[]{2, 0, 2, 1, 1, 0};
        int[] expected = new int[]{0, 0, 1, 1, 2, 2};
        s.sortColors(input);
        Assert.assertTrue(Arrays.equals(input, expected));
    }

    @Test
    public void test2() {
        int[] input = new int[]{2, 0, 1};
        int[] expected = new int[]{0, 1, 2};
        s.sortColors(input);
        Assert.assertTrue(Arrays.equals(input, expected));
    }

    @Test
    public void test3() {
        int[] input = new int[]{2, 0};
        int[] expected = new int[]{0, 2};
        s.sortColors(input);
        Assert.assertTrue(Arrays.equals(input, expected));
    }

    @Test
    public void test4() {
        int[] input = new int[]{2};
        int[] expected = new int[]{2};
        s.sortColors(input);
        Assert.assertTrue(Arrays.equals(input, expected));
    }

    @Test
    public void test5() {
        int[] input = new int[]{2, 2, 2, 1, 1, 1};
        int[] expected = new int[]{1, 1, 1, 2, 2, 2};
        s.sortColors(input);
        Assert.assertTrue(Arrays.equals(input, expected));
    }
}