https://labuladong.gitee.io/algo/4/31/122/

讀完本文，你不僅學會了演算法套路，還可以順便去 LeetCode 上拿下如下題目：

303. 區域和檢索 - 陣列不可變（中等）

304. 二維區域和檢索 - 矩陣不可變（中等）

560. 和為K的子陣列（中等）

———–

字首和技巧適用於快速、頻繁地計算一個索引區間內的元素之和。

一維陣列中的字首和

先看一道例題，力扣第 303 題「區域和檢索 - 陣列不可變」，讓你計算陣列區間內元素的和，這是一道標準的字首和問題：

題目要求你實現這樣一個類：

class NumArray {

    public NumArray(int[] nums) {}
    
    /* 查詢閉區間 [left, right] 的累加和 */
    public int sumRange(int left, int right) {}
}
 

sumRange函式需要計算並返回一個索引區間之內的元素和，沒學過字首和的人可能寫出如下程式碼：

class NumArray {

    private int[] nums;

    public NumArray(int[] nums) {
        this.nums = nums;
    }
    
    public int sumRange(int left, int right) {
        int res = 0;
        for (int i = left; i <= right; i++) {
            res += nums[i];
        }
        return res;
    }
}
 

這樣，可以達到效果，但是效率很差，因為sumRange方法會被頻繁呼叫，而它的時間複雜度是O(N)，其中N代表nums陣列的長度。

這道題的最優解法是使用字首和技巧，將sumRange函式的時間複雜度降為O(1)，說白了就是不要在sumRange裡面用 for 迴圈，咋整？

直接看程式碼實現：

class NumArray {
    // 字首和陣列
    private int[] preSum;

    /* 輸入一個數組，構造字首和 */
    public NumArray(int[] nums) {
        // preSum[0] = 0，便於計算累加和
        preSum = new int[nums.length + 1];
        // 計算 nums 的累加和
        for (int i = 1; i < preSum.length; i++) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
    }
    
    /* 查詢閉區間 [left, right] 的累加和 */
    public int sumRange(int left, int right) {
        return preSum[right + 1] - preSum[left];
    }
}
 

核心思路是我們 new 一個新的陣列preSum出來，preSum[i]記錄nums[0..i-1]的累加和，看圖 10 = 3 + 5 + 2：

看這個preSum陣列，如果我想求索引區間[1, 4]內的所有元素之和，就可以通過preSum[5] - preSum[1]得出。

這樣，sumRange函式僅僅需要做一次減法運算，避免了每次進行 for 迴圈呼叫，最壞時間複雜度為常數O(1)。

這個技巧在生活中運用也挺廣泛的，比方說，你們班上有若干同學，每個同學有一個期末考試的成績（滿分 100 分），那麼請你實現一個 API，輸入任意一個分數段，返回有多少同學的成績在這個分數段內。

那麼，你可以先通過計數排序的方式計算每個分數具體有多少個同學，然後利用字首和技巧來實現分數段查詢的 API：

int[] scores; // 儲存著所有同學的分數
// 試卷滿分 100 分
int[] count = new int[100 + 1]
// 記錄每個分數有幾個同學
for (int score : scores)
    count[score]++
// 構造字首和
for (int i = 1; i < count.length; i++)
    count[i] = count[i] + count[i-1];

// 利用 count 這個字首和陣列進行分數段查詢

接下來，我們看一看字首和思路在實際演算法題中可以如何運用。

二維矩陣中的字首和

這是力扣第 304 題「304. 二維區域和檢索 - 矩陣不可變」，其實和上一題類似，上一題是讓你計運算元陣列的元素之和，這道題讓你計算二維矩陣中子矩陣的元素之和：

比如說輸入的matrix如下圖：

那麼sumRegion([2,1,4,3])就是圖中紅色的子矩陣，你需要返回該子矩陣的元素和 8。

這題的思路和一維陣列中的字首和是非常類似的，如下圖：

如果我想計算紅色的這個子矩陣的元素之和，可以用綠色矩陣減去藍色矩陣減去橙色矩陣最後加上粉色矩陣，而綠藍橙粉這四個矩陣有一個共同的特點，就是左上角就是(0, 0)原點。

那麼我們可以維護一個二維preSum陣列，專門記錄以原點為頂點的矩陣的元素之和，就可以用幾次加減運算算出任何一個子矩陣的元素和：

class NumMatrix {
    // preSum[i][j] 記錄矩陣 [0, 0, i, j] 的元素和
    private int[][] preSum;
    
    public NumMatrix(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length;
        if (m == 0 || n == 0) return;
        // 構造字首和矩陣
        preSum = new int[m + 1][n + 1];
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                // 計算每個矩陣 [0, 0, i, j] 的元素和
                preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] + matrix[i - 1][j - 1] - preSum[i-1][j-1];
            }
        }
    }
    
    // 計運算元矩陣 [x1, y1, x2, y2] 的元素和
    public int sumRegion(int x1, int y1, int x2, int y2) {
        // 目標矩陣之和由四個相鄰矩陣運算獲得
        return preSum[x2+1][y2+1] - preSum[x1][y2+1] - preSum[x2+1][y1] + preSum[x1][y1];
    }
}

這樣，sumRegion函式的複雜度也用字首和技巧優化到了 O(1)。

和為 k 的子陣列

最後聊一道稍微有些困難的字首和題目，力扣第 560 題「和為 K 的子陣列」：

那我把所有子陣列都窮舉出來，算它們的和，看看誰的和等於k不就行了，藉助字首和技巧很容易寫出一個解法：

int subarraySum(int[] nums, int k) {
    int n = nums.length;
    // 構造字首和
    int[] preSum = new int[n + 1];
    preSum[0] = 0; 
    for (int i = 0; i < n; i++)
        preSum[i + 1] = preSum[i] + nums[i];
    
    int res = 0;
    // 窮舉所有子陣列
    for (int i = 1; i <= n; i++)
        for (int j = 0; j < i; j++)
            // 子陣列 nums[j..i-1] 的元素和
            if (preSum[i] - preSum[j] == k)
                res++;

    return res;
}

這個解法的時間複雜度O(N^2)空間複雜度O(N)，並不是最優的解法。不過通過這個解法理解了字首和陣列的工作原理之後，可以使用一些巧妙的辦法把時間複雜度進一步降低。

注意前面的解法有巢狀的 for 迴圈：

for (int i = 1; i <= n; i++)
    for (int j = 0; j < i; j++)
        if (preSum[i] - preSum[j] == k)
            res++;

第二層 for 迴圈在幹嘛呢？翻譯一下就是，在計算，有幾個j能夠使得preSum[i]和preSum[j]的差為k。毎找到一個這樣的j，就把結果加一。

我們可以把 if 語句裡的條件判斷移項，這樣寫：

if (preSum[j] == preSum[i] - k)
    res++;

優化的思路是：我直接記錄下有幾個preSum[j]和preSum[i] - k相等，直接更新結果，就避免了內層的 for 迴圈。我們可以用雜湊表，在記錄字首和的同時記錄該字首和出現的次數。

int subarraySum(int[] nums, int k) {
    int n = nums.length;
    // map：字首和 -> 該字首和出現的次數
    HashMap<Integer, Integer> 
        preSum = new HashMap<>();
    // base case
    preSum.put(0, 1);

    int res = 0, sum0_i = 0;
    for (int i = 0; i < n; i++) {
        sum0_i += nums[i];
        // 這是我們想找的字首和 nums[0..j]
        int sum0_j = sum0_i - k;
        // 如果前面有這個字首和，則直接更新答案
        if (preSum.containsKey(sum0_j))
            res += preSum.get(sum0_j);
        // 把字首和 nums[0..i] 加入並記錄出現次數
        preSum.put(sum0_i, 
            preSum.getOrDefault(sum0_i, 0) + 1);
    }
    return res;
}

比如說下面這個情況，需要字首和 8 就能找到和為k的子陣列了，之前的暴力解法需要遍歷陣列去數有幾個 8，而優化解法藉助雜湊表可以直接得知有幾個字首和為 8。

這樣，就把時間複雜度降到了O(N)，是最優解法了。

字首和技巧就講到這裡，應該說這個演算法技巧是會者不難難者不會，實際運用中還是要多培養自己的思維靈活性，做到一眼看出題目是一個字首和問題。

接下來可閱讀：

• 差分陣列技巧

＿＿＿＿＿＿＿＿＿＿＿＿＿