leetcode_rust/
find_pivot_index.rs1#![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}