leetcode_rust/
find_pivot_index.rs

1#![allow(dead_code)]
2
3pub fn pivot_index(nums: Vec<i32>) -> i32 {
4    if nums.len() < 2 {
5        return nums.len() as i32 - 1;
6    }
7    let mut i = 0;
8    let sum: i32 = nums.iter().sum();
9
10    let mut half = 0;
11    while i < nums.len() {
12        if half * 2 == sum - nums[i] {
13            return i as i32;
14        }
15        half += nums[i];
16        i += 1;
17    }
18    -1
19}
20
21pub fn pivot_index2(nums: Vec<i32>) -> i32 {
22    if nums.len() < 2 {
23        return nums.len() as i32 - 1;
24    }
25    let mut i = 0;
26    let mut right: i32 = nums.iter().sum();
27
28    let mut left = 0;
29    while i < nums.len() {
30        right -= nums[i];
31        if left == right {
32            return i as i32;
33        }
34        left += nums[i];
35        i += 1;
36    }
37    -1
38}
39
40#[cfg(test)]
41mod tests {
42    use super::*;
43
44    #[test]
45    fn test1() {
46        assert_eq!(pivot_index(vec![1, 7, 3, 6, 5, 6]), 3);
47        assert_eq!(pivot_index(vec![1, 2, 3]), -1);
48        assert_eq!(pivot_index(vec![-1, -1, -1, -1, -1, 0]), 2);
49        assert_eq!(pivot_index(vec![-1, -1, -1, 0, 1, 1]), 0);
50        assert_eq!(pivot_index(vec![0]), 0);
51        assert_eq!(pivot_index(vec![1, 0]), 0);
52        assert_eq!(pivot_index(vec![-1, -1, 0, 1, 1, 0]), 5);
53    }
54
55    #[test]
56    fn test2() {
57        assert_eq!(pivot_index2(vec![1, 7, 3, 6, 5, 6]), 3);
58        assert_eq!(pivot_index2(vec![1, 2, 3]), -1);
59        assert_eq!(pivot_index2(vec![-1, -1, -1, -1, -1, 0]), 2);
60        assert_eq!(pivot_index2(vec![-1, -1, -1, 0, 1, 1]), 0);
61        assert_eq!(pivot_index2(vec![0]), 0);
62        assert_eq!(pivot_index2(vec![1, 0]), 0);
63        assert_eq!(pivot_index2(vec![-1, -1, 0, 1, 1, 0]), 5);
64    }
65}