LeetCode-in-All

560. Subarray Sum Equals K

Medium

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.

Example 1:

Input: nums = [1,1,1], k = 2

Output: 2

Example 2:

Input: nums = [1,2,3], k = 3

Output: 2

Constraints:

Solution

import scala.collection.mutable.Map

object Solution {
    def subarraySum(nums: Array[Int], k: Int): Int = {
        var tempSum = 0
        var ret = 0
        val sumCount: Map[Int, Int] = Map(0 -> 1)

        for (i <- nums) {
            tempSum += i

            if (sumCount.contains(tempSum - k)) {
                ret += sumCount(tempSum - k)
            }

            sumCount.updateWith(tempSum) {
                case Some(count) => Some(count + 1)
                case None => Some(1)
            }
        }

        ret
    }
}