LeetCodeKotlin/src/com/boycoder/problems/array/Med_560_SubArrayEqualK.kt at main · chaxiu/LeetCodeKotlin · GitHub

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
package com.boycoder.problems.array

import com.boycoder.utils.asserts

/**
 * @Author: zhutao
 * @datetime: 2021/6/21
 * @desc: https://leetcode-cn.com/problems/subarray-sum-equals-k/
 * Brute force
 * Brute force optimize
 * Prefix sum
 */
object Med_SubArrayEqualK {
    fun subarraySum(nums: IntArray, k: Int): Int {
        return count3(nums, k)
    }

    /**
     * ----------sum[i]---------------
     * ------prefix[j]-------
     *                       ----k----
     * length = i - j
     *
     *           [1,1,1]  k = 2
     * prefix [0, 1,2,3]
     * map   [0 to 1, 1 to 1, 2 to 2, 3 to 3]
     * end = 3, 3 - k = 1, map[1] = 1, count++
     * end = 2, 2 - k = 0, map[0] = 1, count++
     * end = 1, 1 - k = -1, map[-1] not found
     *
     *           [-1, -1, 1] k = 0
     * prefix  [0,-1, -2, -1]
     * map    [0 to 1, -1 to 2, -2 to 1]
     * end = 3, -1 - k = -1, count = 2
     *
     *
     *
     * This solution is similar to [Med_325_MaxSizeSubArray.maxSubArrayLen]
     */
    private fun count3(nums: IntArray, k: Int): Int {
        var count = 0
        var prefix = 0
        val map = hashMapOf<Int, Int>()
        map.put(0 , 1)

        for (num in nums) {
            prefix = prefix + num
            if (map.containsKey(prefix - k)) {
                count  = count + map.getOrDefault(prefix - k, 0)
            }

            map.put(prefix, map.getOrDefault(prefix, 0) + 1)
        }

        return count
    }

    /**
     * ----------sum[i]---------------
     * ------prefix[j]-------
     *                       ----k----
     * loop start end end
     * if (prefix[end] - prefix[start]) count++
     *
     * This solution is similar to [Med_325_MaxSizeSubArray.maxSubArrayLen]
     */
    private fun count2(nums: IntArray, k: Int): Int {
        val prefix = IntArray(nums.size + 1)
        var count = 0

        for (i in 1..nums.size) {
            prefix[i] = prefix[i - 1] + nums[i - 1]
        }

        for (start in 0 until nums.size) {
            for (end in (start+1)..nums.size) {
                if (prefix[end] - prefix[start] == k) {
                    count++
                }
            }
        }

        return count
    }


    /**
     * Brute force + optimize
     */
    private fun count1(nums: IntArray, k: Int): Int {
        var count = 0

        for (start in 0 until nums.size) {
            var sum = 0
            for (end in start until nums.size) {
                sum = sum + nums[end]

                if (sum == k) {
                    count++
                }
            }
        }

        return count
    }

    /**
     * Brute force
     */
    private fun count(nums: IntArray, k: Int): Int {
        var count = 0

        for (start in 0 until nums.size) {

            for (end in start until nums.size) {
                var sum = 0
                for (k in start..end) {
                    sum = sum + nums[k]
                }

                if (sum == k) {
                    count++
                }
            }
        }

        return count
    }
}

fun main() {
    val res = Med_SubArrayEqualK.subarraySum(intArrayOf(1, 1, 1), 2)
    asserts(res, 2)
    val res1 = Med_SubArrayEqualK.subarraySum(intArrayOf(-1, -1, 1), 1)
    asserts(res1, 1)
}