560. Subarray Sum Equals K

Table of Contents

  1. 問題
  2. 雛形
  3. 解き方
  4. コード

問題

560. Subarray Sum Equals K

整数配列 nums と整数 k が与えられます。整数配列 nums の部分配列の和が k となるような部分配列の数を求めてください。

例えば nums = [1,2,3], k = 3 のとき、和が 3 となるような部分配列は [1,2], [3] の 2 つです。

雛形

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
#include <vector>
#include <map>
#include <gtest/gtest.h>
using namespace std;

class Solution
{
public:
    int subarraySum(vector<int> &nums, int k)
    {
        return 0;
    }
};

TEST(TestCase, Case1)
{
    auto s = Solution();
    auto nums = vector<int>{1, 1, 1};
    auto actual = s.subarraySum(nums, 2);
    auto expected = 2;
    EXPECT_EQ(expected, actual);
}
TEST(TestCase, Case2)
{
    auto s = Solution();
    auto nums = vector<int>{1, 2, 3};
    auto actual = s.subarraySum(nums, 3);
    auto expected = 2;
    EXPECT_EQ(expected, actual);
}
TEST(TestCase, Case3)
{
    auto s = Solution();
    auto nums = vector<int>{1, 2, 1, 2, 4, -1};
    auto actual = s.subarraySum(nums, 3);
    auto expected = 4;
    EXPECT_EQ(expected, actual);
}

解き方

累積和を使います。nums の先頭から見ていって総和を求めます。総和 - k の値が過去に現れていたら和が k となるような部分配列があることを示しています。

以下は nums = [1,2,3], k = 3 の例です。nums 配列の添え字を i 、総和を sum とします。

  • sum=0
    • -> 初期化
  • i=0, nums[0]=1, sum=1
    • -> 1-k=-2, 過去に sum に -2 が現れた回数は 0 回
  • i=1, nums[1]=2, sum=3
    • -> 3-k=0, 過去に sumに 0 が現れた回数は 1 回
  • i=2, nums[2]=3, sum=6
    • -> 6-k=3, 過去に sumに 3 が現れた回数は 1 回

コード

総和とその回数を map を使って管理します。一番最初に総和 0 の出現回数を 1 にすることに注意です。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution
{
public:
    int subarraySum(vector<int> &nums, int k)
    {
        int n = nums.size();
        int sum = 0;
        unordered_map<int, int> sum_to_count;
        sum_to_count[0] = 1;
        int answer = 0;
        for (int i = 0; i < n; i++)
        {
            sum += nums[i];
            answer += sum_to_count[sum - k];
            sum_to_count[sum]++;
        }
        return answer;
    }
};