演算法思維 ---- 字首和技巧

阿新 • • 發佈：2020-11-21

字首和技巧

字首和技巧 主要用於解決陣列的子陣列問題。

使用字首和的經典題目：

1. 什麼是字首和

這裡以 LeetCode 的一道題目，560. Subarray Sum Equals K，來講解字首和技巧。

題目中給定一個數組 nums 和一個數字 k，找出陣列中連續的子陣列的和等於給定值 k 的子陣列數量。

這道題就是列舉每一組子陣列，然後求和，統計等於 k 的子陣列數目。但是這樣的操作，顯然時間複雜度很高，而且可能會超時，在面試中面試官明顯不想看到這種程式碼思路。

這裡我們就展示一波 字首和技巧

吧。

字首字首字首，明顯就是當前陣列元素 i 的前面的元素的之和的意思嘛。我們開啟一個數組長度是原陣列長度+1 的新陣列 preSums ，其中每一個元素都是其前面的元素之和，注意不包括當前元素本身。

以 nums = [1, 2, 3] 為例，我們得到的 preSums 就會是 [0, 1, 3, 6]。

你看，字首和陣列就是這麼簡單，當然做題時也要靈活變通，有時候我們不一定需要開啟一個新陣列的，只要一個變數來存放操作過程中的字首和即可。

// 構造字首和陣列
let getPreSum = (nums) => {
  let preSum = [0];
  for (let i = 0; i < nums.length; i++) {
    preSum[i+1] = preSum[i] + nums[i];
  }
  return preSum;
}

既然得到了字首和陣列了，那麼這道題只要字首和之間存在差值為 k 的，就是存在這樣的一個子陣列了。

let subarraySum = (nums, k) => {
  let perSum = getPreSum(nums);
  
  let count = 0;
  for (let i = 1; i <= nums.length; i++) {
    for (let j = 0; j < i; j++) {
      if (preSum[i] - preSum[j] === k) {
        count++;
      }
    }
  }
  return count;
}

上面的解法的時間複雜度是O(n^2)，空間複雜度是O(n)，n 是nums 的長度。到這裡，我們就解決了一道題目了，但是時空複雜度有點高，我們有什麼方法優化呢？

2. 優化

上面的解法中，我們有兩個for迴圈，其主要的作用就是找出當前 i 的字首和，與其前面的字首和元素 j 之間是否存在差值 k，如果存在，那麼元素 j 到 i 的子陣列元素之和就是 k ，顯然這就是一個結果。

因此，我們可以使用一個 Map 結構，在遍歷過程中，如果Map中存在 preSum - k 的鍵值對，說明存在子陣列之和等於 k 的子陣列，如果沒有則把當前 preSum 存入鍵值對。

const subarraySum2 = function (nums, k) {
  let map = new Map();
  map.set(0, 1);

  let result = 0;
  let preSum = 0;
  for (let i = 0; i < nums.length; i++) {
    preSum += nums[i];
    result += map.get(preSum - k) ? map.get(preSum - k) : 0;
    map.set(preSum, map.has(preSum) ? map.get(preSum) + 1 : 1);
  }
  return result;
};

注意一下哦，在開啟字首和陣列時我們的序號0的元素是不存元素的，預設0，因為存在 (k - 0) 的一個字首和；所以建立一個 Map 時，我們也需存入 (0,1) 的一個鍵值對。

至此，我們的時間複雜度降至 O(n) 了。