394. Distribute Elements Into Two Arrays II

The only expensive operation is greaterCount(arr, val) = how many elements already in an array exceed val. Maintain a Binary Indexed Tree (Fenwick tree) per array over compressed values; greaterCount is (size) - (count of values <= val), answered in O(log n).

What's slow, and the fix

The rule for each element depends on greaterCount(arr1, x) vs greaterCount(arr2, x). Recomputing that by scanning is O(n) each, O(n^2) total. Instead keep a Fenwick tree (BIT) for each array, indexed by compressed value, that counts how many elements of that value are present.

Then greaterCount(arr, x) = size(arr) - prefixCount(arr, x), where prefixCount is the number of stored elements with value <= x — a single BIT prefix query in O(log n).

The simulation

1. Coordinate-compress the distinct values so BIT indices are 0..m-1.
2. Seed: arr1 = [nums[0]], arr2 = [nums[1]], updating both BITs.
3. For each remaining x, compute g1 = greaterCount(arr1, x) and g2 = greaterCount(arr2, x) and apply the four-way rule (compare counts; break ties by size; final tie → arr1), updating the chosen BIT.
4. Return arr1 ++ arr2.

Worked example

nums = [2, 1, 3, 3].

• op1: arr1=[2]. op2: arr2=[1].
• op3 (3): greaterCount(arr1,3)=0, greaterCount(arr2,3)=0 → tie; sizes equal (1 and 1) → tie → arr1=[2,3].
• op4 (3): greaterCount(arr1,3)=0, greaterCount(arr2,3)=0 → tie; sizes 2 vs 1 → smaller is arr2 → arr2=[1,3].

result = arr1 ++ arr2 = [2, 3, 1, 3].

Complexity

• Time: O(n log n).
• Space: O(n).

class Solution {
    private int m;
    private int[] ta, tb;

    public int[] resultArray(int[] nums) {
        int n = nums.length;
        int[] srt = Arrays.stream(nums).distinct().sorted().toArray();
        m = srt.length;
        Map<Integer, Integer> idx = new HashMap<>();
        for (int i = 0; i < m; i++) idx.put(srt[i], i);
        ta = new int[m + 1];
        tb = new int[m + 1];

        int[] a = new int[n], b = new int[n];
        int ai = 0, bi = 0;
        a[ai++] = nums[0];
        update(ta, idx.get(nums[0]));
        b[bi++] = nums[1];
        update(tb, idx.get(nums[1]));

        for (int k = 2; k < n; k++) {
            int x = nums[k];
            int i = idx.get(x);
            int gA = ai - query(ta, i);  // elements in a strictly greater than x
            int gB = bi - query(tb, i);
            if (gA > gB || (gA == gB && ai <= bi)) {
                a[ai++] = x;
                update(ta, i);
            } else {
                b[bi++] = x;
                update(tb, i);
            }
        }
        int[] res = new int[n];
        System.arraycopy(a, 0, res, 0, ai);
        System.arraycopy(b, 0, res, ai, bi);
        return res;
    }

    private void update(int[] t, int i) {
        for (i += 1; i <= m; i += i & (-i)) t[i] += 1;
    }

    private int query(int[] t, int i) {  // count of indices <= i
        int r = 0;
        for (i += 1; i > 0; i -= i & (-i)) r += t[i];
        return r;
    }
}