Smallest subarrays

You are given two arrays $A$ and $B$ of length $N$.

For every index $i$ in array $A$, find the length of the smallest subarray starting from index $i$ such that there exist at least $B[i]$ elements with a value greater than or equal to $A[i]$ in the subarray. If no such subarray exists, then print $-1$.

Example:

Input:  A[] = { 8, 2, 4, 1, 9 }, B[] = { 2, 3, 5, 1, 1 }
Output: 5 4 -1 1 1 

Approach

Java

import java.util.Arrays;
import java.util.Comparator;

public class SmallestSubarrays {
    public static void main(String[] args) throws Exception {

        int n = 5;
        int A[] = { 8, 2, 4, 1, 9 };
        int B[] = { 2, 3, 5, 1, 1 };
        int[][] a = new int[n][4];

        for (int i = 0; i < n; ++i) {
            a[i][0] = i;
            a[i][1] = A[i];
        }
        for (int i = 0; i < n; ++i)
            a[i][2] = B[i];

        Arrays.sort(a, Comparator.comparingInt(o -> o[1]));

        int[] tree = new int[4 * n];

        for (int i = 0; i < n; ++i) {
            int ci = a[i][0];
            int fw = a[i][2];

            int missing = query(tree, 0, n - 1, ci, n - 1, 1);
            int poss = n - ci;

            if (poss - missing < fw)
                a[i][3] = -1;
            else
                a[i][3] = solve(tree, ci, fw, n);

            update(tree, 0, n - 1, ci, 1);
        }

        Arrays.sort(a, Comparator.comparingInt(o -> o[0]));

        for (int i = 0; i < n; ++i)
            System.out.print(a[i][3] + " ");
    }

    static void update(int[] tree, int beg, int end, int idx, int pos) {
        if (beg == end && beg == idx) {
            tree[pos] = 1;
            return;
        }

        int mid = (beg + end) / 2;
        if (idx <= mid)
            update(tree, beg, mid, idx, 2 * pos);
        else
            update(tree, mid + 1, end, idx, 2 * pos + 1);
        tree[pos] = tree[2 * pos] + tree[2 * pos + 1];
    }

    static int query(int[] tree, int beg, int end, int l, int r, int pos) {
        if (r < beg || end < l)
            return 0;
        else if (l <= beg && end <= r)
            return tree[pos];

        int mid = (beg + end) / 2;
        int left = query(tree, beg, mid, l, r, 2 * pos);
        int right = query(tree, mid + 1, end, l, r, 2 * pos + 1);
        return left + right;
    }

    static int solve(int[] tree, int ci, int fw, int n) {
        int beg = ci;
        int end = n - 1;
        int mid = (beg + end) / 2;
        int ans = -1;
        while (beg <= end) {
            mid = (beg + end) / 2;
            int missing = query(tree, 0, n - 1, ci, mid, 1);
            if (fw + missing == mid - ci + 1) {
                ans = fw + missing;
                end = mid - 1;
            } else if (fw + missing >= mid - ci + 1)
                beg = mid + 1;
            else
                end = mid - 1;
        }
        return ans;
    }

}