寻找数组的中心下标 Find Pivot Index

问题描述

这个问题的本质是要理解 pivot index 的数学含义:

其关系如上图所示, 从中我们可以发现 pivot index 存在时要满足这些条件:

• pivot 元素将数组分成左右两部分, 分别叫 prefix array 和 suffix array
• sum(prefix array) = sum(suffix array)

然而, 还有一个隐藏的条件是:

sum(prefix array) + pivot + sum(suffix array) = sum(array)

Brute force

上面的公式还可以转换成:

sum(prefix array) + pivot + sum(prefix array) = sum(array)

从左到右遍历数组中的所有元素, 循环终止的条件就是找到:

sum(prefix array) + pivot + sum(prefix array) = sum(array)

以下是代码实现:

#![allow(unused)]
fn main() {
// Brute force
pub fn pivot_index1(nums: Vec<i32>) -> i32 {
    // 第一步: 计算数组中所有元素的和
    let sum: i32 = nums.iter().sum();

    // 第二步: 从左侧遍历数组, 并计算数组的前缀和 prefix sum
    // 如果前缀和 * 2 + 当前的元组, 其和等于所有元素之和,
    // 那么当前元素所在位置就是 pivot index.
    let mut prefix_sum: i32 = 0;
    for (index, num) in nums.iter().enumerate() {
        if prefix_sum * 2 + num == sum {
            return index as i32;
        }
        prefix_sum += num;
    }
    -1
}
}

这种算法的特点是:

• 时间复杂度是 O(n)
• 空间复杂度是 O(1)

前缀和 Prefix Sum

直接计算 prefix sum 和 suffix sum.

这个方法跟上面的类似, 但是更好理解一些, 它不需要最开始说的公式转换:

• 首先计算所有元素的和, 作为 suffix sum; 同时将 prefix sum 初始化为 0
• 然后从左到右遍历数组, 将该本元从 suffix sum 中减去
  • 如果此时 prefix sum == suffix sum, 则当前元素就是 pivot, 当前位置就是 pivot index, 直接返回
  • 否则将该元素加到 prefix sum 中

算法实现如下所示:

#![allow(unused)]
fn main() {
// Prefix Sum
// 直接计算 prefix sum 和 suffix sum
pub fn pivot_index2(nums: Vec<i32>) -> i32 {
    // 第一步: 计算数组中所有元素的和, 并作为 suffix sum
    let mut suffix_sum: i32 = nums.iter().sum();

    // 第二步: 从左侧遍历数组, 并调整 prefix_sum 和 suffix_sum 的值
    // 如果它们相等了, 就终止循环
    let mut prefix_sum: i32 = 0;
    for (index, num) in nums.iter().enumerate() {
        suffix_sum -= num;
        if prefix_sum == suffix_sum {
            return index as i32;
        }
        prefix_sum += num;
    }
    -1
}
}

这种算法的特点是:

• 时间复杂度是 O(n)
• 空间复杂度是 O(1)