algorithm/leetcode/BinarySubarraysWithSum.java at master · skyczhao/algorithm · 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
package leetcode;


/**
 * BinarySubarraysWithSum
 * https://leetcode-cn.com/problems/binary-subarrays-with-sum/
 * 930. 和相同的二元子数组
 * https://leetcode-cn.com/problems/binary-subarrays-with-sum/solution/ya-suo-zu-he-jie-ti-by-oshdyr-apbq/
 *
 * @author tobin
 * @since 2021-07-08
 */
public class BinarySubarraysWithSum {

    public static void main(String[] args) {
        BinarySubarraysWithSum sol = new BinarySubarraysWithSum();
        System.out.println(sol.numSubarraysWithSum(new int[]{1, 0, 1, 0, 1}, 2));
        System.out.println(sol.numSubarraysWithSum(new int[]{0, 0, 0, 0, 0}, 0));
    }

    public int numSubarraysWithSum(int[] nums, int goal) {
        int oneSize = 0;
        int[] zeroBefore = new int[nums.length];
        int zeroCount = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] > 0) {
                zeroBefore[oneSize] = zeroCount;
                oneSize++;

                zeroCount = 0;
            } else {
                zeroCount++;
            }
        }
        int zeroFinal = zeroCount;

        int times = 0;
        if (goal == 0) {
            for (int i = 0; i < oneSize; i++) {
                times += (zeroBefore[i] * (zeroBefore[i] + 1) / 2);
            }
            times += (zeroFinal * (zeroFinal + 1) / 2);
        } else {
            for (int i = 0; i < oneSize; i++) {
                int e = i + goal;
                if (e > oneSize) {
                    break;
                }

                int zeroHead = zeroBefore[i];
                int zeroTail = zeroFinal;
                if (e < oneSize) {
                    zeroTail = zeroBefore[e];
                }
//                if (zeroHead == 0) {
//                    zeroHead = 1;
//                }
                zeroHead += 1;
//                if (zeroTail == 0) {
//                    zeroTail = 1;
//                }
                zeroTail += 1;
                times += (zeroHead * zeroTail);
            }
        }
        return times;
    }

//    public int numSubarraysWithSum(int[] nums, int goal) {
//
//        int[] accumulator = new int[nums.length];
//        for (int i = 0; i < nums.length; i++) {
//            if (i > 0) {
//                accumulator[i] = accumulator[i - 1] + nums[i];
//            } else {
//                accumulator[0] = nums[0];
//            }
//        }
//
//        int times = 0;
//        for (int i = 0; i < nums.length; i++) {
//            for (int j = i; j < nums.length; j++) {
//                int sum = accumulator[j];
//                if (i > 0) {
//                    sum = sum - accumulator[i - 1];
//                }
//                if (sum == goal) {
//                    times++;
//                }
//            }
//        }
//
//        return times;
//    }
}