GeeksForGeeks-Problem/15 January Count the Substring at main · Michael-Calce/GeeksForGeeks-Problem · 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
class Solution{
public:
  void merging(vector<long long> &pre, int l, int m, int h, long long &inv) {
        int i = l, j = m + 1, k = 0;
        vector<long long> temp(h - l + 1);

        while(i <= m and j <= h) {
            if(pre[i] <= pre[j])
                temp[k++] = pre[i++];
            else {
                inv += (m - i + 1);
                temp[k++] = pre[j++];
            }
        }

        while(i <= m) temp[k++] = pre[i++];

        while(j <= h) temp[k++] = pre[j++];

        k = 0;
        for(int c = l; c <= h; c++)
            pre[c] = temp[k++];
    }

    void mergeSort(vector<long long> &pre, int l, int h, long long &inv) {
        if(l < h) {
            int mid = (l + h) >> 1;
            mergeSort(pre, l, mid, inv);
            mergeSort(pre, mid + 1, h, inv);
            merging(pre, l, mid, h, inv);
        }
    }

    long long inversions(vector<long long> &pre) {
        int n = pre.size();
        long long inv = 0;

        mergeSort(pre, 0, n - 1, inv);

        return inv;
    }

    long long countSubstring(string s){
        int n = s.size();

        vector<int> nums(n);
        for(int i = 0; i < n; i++) {
            nums[i] = s[i] - '0';
            if(nums[i] == 0)
                nums[i] = -1;
        }

        vector<long long> pre_sum(n);
        long long sum = 0;
        for(int i = 0; i < n; i++) {
            sum += nums[i];
            pre_sum[i] = sum;
        }

        long long count = 0; // to store valid substrings;
        for(int i = 0; i < n; i++) {
            if(pre_sum[i] > 0) count++;
        }

        reverse(pre_sum.begin(), pre_sum.end());

        return count + inversions(pre_sum);
    }
};