🔗 🟢 724. Find Pivot Index

tags: `陣列(Array)` `前綴和(Prefix Sum)`

題意

給定一個整數陣列 $nums$ ，請計算這個陣列的 中心索引(Pivot Index) 。

中心索引(Pivot Index) 是陣列的一個索引，其左側所有元素相加的和等於右側所有元素相加的和。

如果中心索引位於陣列的最左側，那麼左側數之和視為 $0$ ，因為沒有元素在索引的左側，如果位於陣列的最右側同理。

如果陣列有多個中心索引，則返回最左側的那一個，如果陣列不存在中心索引，返回 $-1$ 。

思路：前綴和(Prefix Sum)

根據題意，若存在中心索引 $i$ ，則有 $nums[0] + nums[1] + \cdots + nums[i-1] = nums[i+1] + nums[i+2] + \cdots + nums[n-1]$ ，其中 $n$ 為陣列長度。令 $tot$ 為所有元素的總和， $s$ 為當前索引 $i$ 左側所有元素的和，則有 $s + nums[i] + s = tot$ ，即 $2 \times s + nums[i] = tot$ 。

由於 $s = \displaystyle\sum_{j=0}^{i-1} nums[j]$ ，因此可以使用前綴和(Prefix Sum) 的概念來解決這個問題。在遍歷陣列的過程中，維護一個前綴和 $s$ ，對於每個索引 $i$ ，檢查 $2 \times s + nums[i] == tot$ 是否成立。如果成立，則 $i$ 就是中心索引；否則則累加 $s$ 並繼續遍歷。最後，如果不存在這樣的索引，則返回 $-1$ 。

複雜度分析

時間複雜度: $\mathcal{O}(n)$ ，其中 $n$ 為陣列的長度。
空間複雜度: $\mathcal{O}(1)$ 。

class Solution:
    def pivotIndex(self, nums: List[int]) -> int:
        tot = sum(nums)
        s = 0 # prefix sum
        for i, x in enumerate(nums):
            if 2 * s + x == tot:
                return i
            s += x
        return -1

class Solution {
public:
    int pivotIndex(vector<int>& nums) {
        int n = nums.size();
        int tot = accumulate(nums.begin(), nums.end(), 0);
        int s = 0;  // prefix sum
        for (int i = 0; i < n; ++i) {
            if ((s << 1) + nums[i] == tot) return i;
            s += nums[i];
        }
        return -1;
    }
};