75. Sort Colors Medium

1/**
2 * [75] Sort Colors
3 *
4 * Given an array nums with n objects colored red, white, or blue, sort them <a href="https://en.wikipedia.org/wiki/In-place_algorithm" target="_blank">in-place</a> so that objects of the same color are adjacent, with the colors in the order red, white, and blue.
5 * We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.
6 * You must solve this problem without using the library's sort function.
7 *  
8 * Example 1:
9 * 
10 * Input: nums = [2,0,2,1,1,0]
11 * Output: [0,0,1,1,2,2]
12 * 
13 * Example 2:
14 * 
15 * Input: nums = [2,0,1]
16 * Output: [0,1,2]
17 * 
18 *  
19 * Constraints:
20 * 
21 * n == nums.length
22 * 1 <= n <= 300
23 * nums[i] is either 0, 1, or 2.
24 * 
25 *  
26 * Follow up: Could you come up with a one-pass algorithm using only constant extra space?
27 * 
28 */
29pub struct Solution {}
30
31// problem: https://leetcode.com/problems/sort-colors/
32// discuss: https://leetcode.com/problems/sort-colors/discuss/?currentPage=1&orderBy=most_votes&query=
33
34// submission codes start here
35
36impl Solution {
37    pub fn sort_colors(nums: &mut Vec<i32>) {
38        let (mut zero, mut two, mut i) = (0, nums.len() - 1, 0);
39        while i <= two {
40            if nums[i] == 2 {
41                Self::swap(nums, i, two);
42                if two == 0 {
43                    return
44                }
45                two -= 1;
46            } else if nums[i] == 0 {
47                Self::swap(nums, i, zero);
48                zero += 1;
49                i += 1;
50            } else {
51                i += 1;
52            }
53        }
54    }
55
56    fn swap(nums: &mut Vec<i32>, x: usize, y: usize) {
57        let tmp = nums[x];
58        nums[x] = nums[y];
59        nums[y] = tmp;
60    }
61}
62
63// submission codes end
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_75() {
71        let mut i = vec![1,2,1,1,0];
72        Solution::sort_colors(&mut i);
73        assert_eq!(i, vec![0,1,1,1,2]);
74    }
75
76    #[test]
77    fn test_75_1() {
78        let mut i = vec![2,0,1];
79        Solution::sort_colors(&mut i);
80        assert_eq!(i, vec![0,1,2]);
81    }
82
83    #[test]
84    fn test_75_2() {
85        let mut i = vec![2];
86        Solution::sort_colors(&mut i);
87        assert_eq!(i, vec![2]);
88    }
89
90    #[test]
91    fn test_75_3() {
92        let mut i = vec![2,2];
93        Solution::sort_colors(&mut i);
94        assert_eq!(i, vec![2,2]);
95    }
96}
97