-
Notifications
You must be signed in to change notification settings - Fork 4
/
main.rs
78 lines (72 loc) · 2.04 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
fn main() {
let mut input = vec![2,0,2,1,1,0];
Solution::sort_colors(&mut input);
}
struct Solution {}
impl Solution {
/// [one..=two] (both inclusive) are unordered
/// three pointers
pub fn sort_colors(nums: &mut Vec<i32>) {
if nums.len() <= 1 { return }
let (mut zero, mut two) = (0, nums.len() - 1);
while nums[two] == 2 {
if let Some(t) = two.checked_sub(1) {
two = t;
} else { return }
}
while zero < nums.len() && nums[zero] == 0 {
zero += 1;
}
let mut one = zero;
while one <= two {
if one == two {
if nums[one] == 0 { nums.swap(zero, one) }
return
}
if nums[one] == 2 {
nums.swap(one, two);
two -= 1; // two cannot be 0 cause one < two
} else if nums[one] == 0 {
nums.swap(zero, one);
zero += 1;
one += 1;
} else {
one += 1;
}
}
}
}
#[cfg(test)]
mod test {
use crate::*;
#[test]
fn basic() {
let mut input = vec![2,0,2,1,1,0];
Solution::sort_colors(&mut input);
let output = vec![0,0,1,1,2,2];
assert_eq!(input, output);
}
#[test]
fn edge() {
let mut input = vec![2,2];
Solution::sort_colors(&mut input);
let output = vec![2,2];
assert_eq!(input, output);
let mut input = vec![1,1];
Solution::sort_colors(&mut input);
let output = vec![1,1];
assert_eq!(input, output);
let mut input = vec![0,0];
Solution::sort_colors(&mut input);
let output = vec![0,0];
assert_eq!(input, output);
let mut input = vec![0];
Solution::sort_colors(&mut input);
let output = vec![0];
assert_eq!(input, output);
let mut input = vec![];
Solution::sort_colors(&mut input);
let output = vec![];
assert_eq!(input, output);
}
}