Prefix Sum 기법 배워보기

Prefix sum

• Prefix sum 기법은 배열의 부분 합을 미리 구해놓고 계산의 횟수를 줄이는 기법이다.
• ex) [1, 2, 3, 4, 5] 이라는 배열이 있다면, 각 부분합인 [1, 3, 6, 10, 15] 를 저장해둔다.
• Prefix sum 을 미리 구해둔 상태에서 아래와 같은 계산을 쉽게 할 수 있다.
  • 1~5 까지의 합 - 1~3 까지의 합 을 구하라고 한다면, 각각을 다시 반복문으로 반복할 필요 없이 15 - 6 을 구하면 된다. (15가 15까지의 합, 6이 13까지의 합이기 때문)
  • 원본 배열에서 인덱스 2번부터 4번까지의 합(sum(array[2], array[3], array[4]))을 구하라고 한다면, prefixSum[4] - prefixSum[2] 로 쉽게 구할 수 있다.
• 하나의 배열에서 계산이 많으며 배열의 길이가 길다면, 다시 계산하는 연산의 수가 줄어드므로 Prefix sum 의 효율은 보통 좋아진다.
• 시간복잡도 O(n) 이후에는 합계와 관련된 문제를 O(1) 에 해결할 수 있다.

연습문제 1

• queries 는 배열의 인덱스 시작과 끝 범위를 가지고 있다.
• queries 가 가리키는 배열의 범위 숫자의 총 합이 limit 보다 작으면 true 아니면 false 를 반환하라.

입력

• nums: 정수 배열
  • ex) [1, 6, 3, 2, 7, 2]
• queries: 배열 인덱스 범위를 원소로 하는 배열
  • ex) [[0, 3], [2, 5], [2, 4]]
• limit: 정수
  • ex) 13

출력

• [true, false, true]

내가 짠 코드

function ps(nums, queries, limit) {
  // prefix sum 에 해당하는 임시 배열을 만든다.
  const prefixSum = [nums[0]];
  nums.slice(1).forEach((e, i) => prefixSum.push(e + prefixSum[i]));

  return queries.map(
    ([l, r]) => prefixSum[r] - (prefixSum[l] - nums[l]) < limit
  );
}

ps(
  [1, 6, 3, 2, 7, 2],
  [
    [0, 3],
    [2, 5],
    [2, 4],
  ],
  13
);

모범 답안

/**
 * @param {number[]} nums
 * @param {number[][]} queries
 * @param {number} limit
 * @return {boolean[]}
 */
var answerQueries = function (nums, queries, limit) {
  let prefix = [nums[0]];
  for (let i = 1; i < nums.length; i++) {
    prefix.push(nums[i] + prefix[prefix.length - 1]);
  }

  let ans = [];
  for (const [x, y] of queries) {
    let curr = prefix[y] - prefix[x] + nums[x];
    ans.push(curr < limit);
  }

  return ans;
};

answerQueries(
  [1, 6, 3, 2, 7, 2],
  [
    [0, 3],
    [2, 5],
    [2, 4],
  ],
  13
);

연습문제 2

https://leetcode.com/problems/number-of-ways-to-split-array/

내가 짠 코드

function ps2(nums) {
  const p = [nums[0]];
  nums.slice(1).forEach((e, i) => p.push(e + p[i]));

  let ans = 0;
  p.slice(0, -1).forEach((e, i) => {
    const r = p[p.length - 1];
    if (e >= r - e) ans++;
  });

  return ans;
}

ps2([10, 4, -8, 7]);

모범 답안

/**
 * @param {number[]} nums
 * @return {number}
 */
var waysToSplitArray = function (nums) {
  let ans = 0,
    leftSection = 0,
    total = 0;
  for (const num of nums) {
    total += num;
  }

  for (let i = 0; i < nums.length - 1; i++) {
    leftSection += nums[i];
    let rightSection = total - leftSection;
    if (leftSection >= rightSection) {
      ans++;
    }
  }

  return ans;
};

우리는 실제로 마지막 숫자만 이용하고 나머지 숫자는 순서대로 이용했으므로, 공간복잡도를 약간 줄일 수 있다.