Count of substrings in a Binary String that contains more 1s than 0s

import java.util.Arrays;

public class SubstringsCount {

    public static void merge(int[] arr, int left, int mid, int right, int[] inversions) {
        // Merge two sorted partitions and count inversions
        int[] temp = new int[right - left + 1];
        int i = left;
        int j = mid + 1;
        int k = 0;

        while (i <= mid && j <= right) {
            if (arr[i] <= arr[j]) {
                temp[k] = arr[i];
                k++;
                i++;
            } else {
                // Counting inversions when an element from the right partition is chosen
                inversions[0] += (mid - i + 1);
                temp[k] = arr[j];
                k++;
                j++;
            }
        }

        while (i <= mid) {
            temp[k] = arr[i];
            k++;
            i++;
        }

        while (j <= right) {
            temp[k] = arr[j];
            k++;
            j++;
        }

        k = 0;
        for (int a = left; a <= right; a++) {
            arr[a] = temp[k];
            k++;
        }
    }

    public static void mergeSort(int[] arr, int left, int right, int[] inversions) {
        // Recursive merge sort algorithm
        if (left < right) {
            int mid = (left + right) / 2;
            mergeSort(arr, left, mid, inversions);
            mergeSort(arr, mid + 1, right, inversions);
            merge(arr, left, mid, right, inversions);
        }
    }

    public static int countInversions(int[] arr) {
        // Calculate the number of inversions using the modified merge sort
        int[] inversions = {0};
        mergeSort(arr, 0, arr.length - 1, inversions);
        return inversions[0];
    }

    public static int getSubsCount(String inputStr) {
        int n = inputStr.length();
        int[] nums = new int[n];

        // Convert binary string to an array of 1s and -1s
        for (int i = 0; i < n; i++) {
            nums[i] = Character.getNumericValue(inputStr.charAt(i));
            if (nums[i] == 0) {
                nums[i] = -1;
            }
        }

        int[] pref = new int[n];
        int s = 0;

        // Calculate the prefix sum array
        for (int i = 0; i < n; i++) {
            s += nums[i];
            pref[i] = s;
        }

        // Count the number of valid substrings
        int cnt = (int) Arrays.stream(pref).filter(x -> x > 0).count();

        // Reverse the prefix sum array
        for (int i = 0; i < n / 2; i++) {
            int temp = pref[i];
            pref[i] = pref[n - 1 - i];
            pref[n - 1 - i] = temp;
        }

        // Calculate the number of inversions in the reversed prefix sum array
        int inversions = countInversions(pref);

        // Calculate the final result by adding counts of valid substrings and inversions
        return cnt + inversions;
    }

    // Driver Code
    public static void main(String[] args) {
        // Given Input
        String inputStr = "10011101";

        // Function Call
        int ans = getSubsCount(inputStr);

        // Print the result
        System.out.println(ans);
    }
}

import java.util.*;

class Main {
    static List<Integer> tree;

    static void update(int start, int end, int parent, long index) {
        if (start > end) {
            return;
        }
        if (start == end) {
            tree.set(parent, tree.get(parent) + 1);
            return;
        }
        int mid = (start + end) / 2;
        if (index > mid) {
            update(mid + 1, end, 2 * parent + 2, index);
        }
        else {
            update(start, mid, 2 * parent + 1, index);
        }
        tree.set(parent, tree.get(2 * parent + 1) + tree.get(2 * parent + 2));
    }

    static int query(int start, int end, int parent, int qstart, int qend) {
        if (qstart > end || qend < start) {
            return 0;
        }
        if (qstart <= start && qend >= end) {
            return tree.get(parent);
        }
        int mid = (start + end) / 2;
        int L = query(start, mid, 2 * parent + 1, qstart, qend);
        int R = query(mid + 1, end, 2 * parent + 2, qstart, qend);
        return L + R;
    }

    static int getSubsCount(String S) {
        int n = S.length();
        tree = new ArrayList<Integer>(Collections.nCopies(4 * 2 * n + 1, 0));

        int shift = n;
        long currSum = 0;
        long cnt = 0;

        update(0, 2 * n, 0, 0 + shift);

        for (int i = 0; i < n; i++) {
            currSum += (S.charAt(i) == '1' ? 1 : -1);

            /*    prefix[j]-prefix[i]>=1
                => prefix[i]<=1-prefix[j]
            */

            int lessThan = (int) (currSum + shift) - 1;
            cnt += query(0, 2 * n, 0, 0, lessThan);

            update(0, 2 * n, 0, currSum + shift);
        }
        return (int) cnt;
    }

    public static void main(String[] args) {
        String input = "10011101";
        int ans = getSubsCount(input);
        System.out.println(ans);
    }
}