Distribute Elements Into Two Arrays II

Hard

#2728Time: O(n^2). The main loop runs `n-2` times. In each iteration `i`, calculating `greaterCount` takes `O(size of arr1)` and `O(size of arr2)`, which are both at most `O(i)`. The total time is the sum of `O(i)` for `i` from 1 to `n-1`, which is `O(n^2)`.Space: O(n). We use two lists, `arr1` and `arr2`, to store the `n` elements of the input array.1 company

Data structures

Array Binary Indexed Tree Segment Tree

Companies

Autodesk

Prompt

You are given a 1-indexed array of integers nums of length n.

We define a function greaterCount such that greaterCount(arr, val) returns the number of elements in arr that are strictly greater than val.

You need to distribute all the elements of nums between two arrays arr1 and arr2 using n operations. In the first operation, append nums[1] to arr1. In the second operation, append nums[2] to arr2. Afterwards, in the i^th operation:

If greaterCount(arr1, nums[i]) > greaterCount(arr2, nums[i]), append nums[i] to arr1.
If greaterCount(arr1, nums[i]) < greaterCount(arr2, nums[i]), append nums[i] to arr2.
If greaterCount(arr1, nums[i]) == greaterCount(arr2, nums[i]), append nums[i] to the array with a lesser number of elements.
If there is still a tie, append nums[i] to arr1.

The array result is formed by concatenating the arrays arr1 and arr2. For example, if arr1 == [1,2,3] and arr2 == [4,5,6], then result = [1,2,3,4,5,6].

Return the integer array result.

Example 1:

Input: nums = [2,1,3,3]
Output: [2,3,1,3]
Explanation: After the first 2 operations, arr1 = [2] and arr2 = [1].
In the 3rd operation, the number of elements greater than 3 is zero in both arrays. Also, the lengths are equal, hence, append nums[3] to arr1.
In the 4th operation, the number of elements greater than 3 is zero in both arrays. As the length of arr2 is lesser, hence, append nums[4] to arr2.
After 4 operations, arr1 = [2,3] and arr2 = [1,3].
Hence, the array result formed by concatenation is [2,3,1,3].

Example 2:

Input: nums = [5,14,3,1,2]
Output: [5,3,1,2,14]
Explanation: After the first 2 operations, arr1 = [5] and arr2 = [14].
In the 3rd operation, the number of elements greater than 3 is one in both arrays. Also, the lengths are equal, hence, append nums[3] to arr1.
In the 4th operation, the number of elements greater than 1 is greater in arr1 than arr2 (2 > 1). Hence, append nums[4] to arr1.
In the 5th operation, the number of elements greater than 2 is greater in arr1 than arr2 (2 > 1). Hence, append nums[5] to arr1.
After 5 operations, arr1 = [5,3,1,2] and arr2 = [14].
Hence, the array result formed by concatenation is [5,3,1,2,14].

Example 3:

Input: nums = [3,3,3,3]
Output: [3,3,3,3]
Explanation: At the end of 4 operations, arr1 = [3,3] and arr2 = [3,3].
Hence, the array result formed by concatenation is [3,3,3,3].

Constraints:

3 <= n <= 10⁵
1 <= nums[i] <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach directly simulates the process described in the problem. For each element nums[i], we iterate through the current arr1 and arr2 to compute greaterCount. This is straightforward to implement but inefficient for large inputs.

Algorithm

1. Initialize arr1 with nums[0] and arr2 with nums[1].
1. Iterate i from 2 to n-1:
1. Let val = nums[i].
1. Calculate count1 by iterating through arr1 and counting elements > val.
1. Calculate count2 by iterating through arr2 and counting elements > val.
1. If count1 > count2, append val to arr1.
1. Else if count1 < count2, append val to arr2.
1. Else (if counts are equal):
1. If `arr1.size() <= arr2.size()`, append `val` to `arr1`.
1. Else, append val to arr2.
1. Concatenate arr1 and arr2 into a result array.

Walkthrough

We initialize two dynamic arrays, arr1 and arr2. We place nums[0] in arr1 and nums[1] in arr2 as per the problem statement. Then, we loop through the rest of the nums array from the third element (i=2). In each iteration, we calculate greaterCount(arr1, nums[i]) and greaterCount(arr2, nums[i]) by performing a linear scan over arr1 and arr2. Based on the comparison of these counts and the sizes of arr1 and arr2, we append nums[i] to the appropriate array. The greaterCount function simply loops through an array and increments a counter for each element that is strictly greater than the given value. After processing all elements, we concatenate arr1 and arr2 to get the final result.

1import java.util.ArrayList;2import java.util.List;3 4class Solution {5    private int greaterCount(List<Integer> arr, int val) {6        int count = 0;7        for (int num : arr) {8            if (num > val) {9                count++;10            }11        }12        return count;13    }14 15    public int[] resultArray(int[] nums) {16        if (nums.length <= 2) {17            return nums;18        }19 20        List<Integer> arr1 = new ArrayList<>();21        List<Integer> arr2 = new ArrayList<>();22 23        arr1.add(nums[0]);24        arr2.add(nums[1]);25 26        for (int i = 2; i < nums.length; i++) {27            int val = nums[i];28            int count1 = greaterCount(arr1, val);29            int count2 = greaterCount(arr2, val);30 31            if (count1 > count2) {32                arr1.add(val);33            } else if (count1 < count2) {34                arr2.add(val);35            } else {36                if (arr1.size() <= arr2.size()) {37                    arr1.add(val);38                } else {39                    arr2.add(val);40                }41            }42        }43 44        int[] result = new int[nums.length];45        int index = 0;46        for (int num : arr1) {47            result[index++] = num;48        }49        for (int num : arr2) {50            result[index++] = num;51        }52        return result;53    }54}

Complexity

Time

O(n^2). The main loop runs `n-2` times. In each iteration `i`, calculating `greaterCount` takes `O(size of arr1)` and `O(size of arr2)`, which are both at most `O(i)`. The total time is the sum of `O(i)` for `i` from 1 to `n-1`, which is `O(n^2)`.

Space

O(n). We use two lists, `arr1` and `arr2`, to store the `n` elements of the input array.

Trade-offs

Pros

Simple to understand and implement.
Follows the problem description directly.

Cons

Inefficient due to repeated linear scans.
Will result in Time Limit Exceeded (TLE) for large inputs as specified in the constraints (n <= 10^5).

To overcome the O(n^2) complexity of the naive approach, we need a faster way to compute greaterCount. A Fenwick Tree (or Binary Indexed Tree, BIT) is a data structure that can efficiently find the number of elements in a range and handle point updates. Since the values in nums can be large, we first use coordinate compression to map them to a smaller range of indices suitable for the BIT.

Algorithm

1. Create a sorted list of unique elements from nums and a map valueToRank to store the 1-based rank of each unique value. Let m be the number of unique values.
1. Initialize two Fenwick Trees, bit1 and bit2, of size m.
1. Initialize two lists, arr1 and arr2.
1. Add nums[0] to arr1 and update bit1 at its rank: bit1.update(valueToRank.get(nums[0]), 1).
1. Add nums[1] to arr2 and update bit2 at its rank: bit2.update(valueToRank.get(nums[1]), 1).
1. Iterate i from 2 to n-1:
1. Let val = nums[i] and rank = valueToRank.get(val).
1. Calculate count1 = arr1.size() - bit1.query(rank).
1. Calculate count2 = arr2.size() - bit2.query(rank).
1. Apply the distribution rules based on count1, count2, and list sizes.
1. If val is added to arr1, call bit1.update(rank, 1).
1. If val is added to arr2, call bit2.update(rank, 1).
1. Concatenate arr1 and arr2 to form the final result.

Walkthrough

The main bottleneck is the greaterCount function. We can optimize this query from O(n) to O(log n) using a suitable data structure.

Coordinate Compression: The values in nums can be up to 10^9, which is too large for direct indexing. We first find all unique values in nums, sort them, and create a mapping from each value to its rank (1-indexed position in the sorted unique list). This compresses the value range to [1, m], where m is the number of unique elements (m <= n).

Fenwick Tree (BIT): We use two Fenwick Trees, bit1 and bit2, corresponding to arr1 and arr2. Each BIT will store the frequencies of the ranks of numbers present in its corresponding array. The size of each BIT will be m+1.

This approach reduces the time for each greaterCount query to O(log m) and each update to O(log m), leading to an overall O(n log n) solution.

1import java.util.*;2 3class FenwickTree {4    private int[] bit;5    private int size;6 7    public FenwickTree(int size) {8        this.size = size;9        this.bit = new int[size + 1];10    }11 12    public void update(int index, int delta) {13        while (index <= size) {14            bit[index] += delta;15            index += index & -index;16        }17    }18 19    public int query(int index) {20        int sum = 0;21        while (index > 0) {22            sum += bit[index];23            index -= index & -index;24        }25        return sum;26    }27}28 29class Solution {30    public int[] resultArray(int[] nums) {31        int n = nums.length;32        if (n <= 2) {33            return nums;34        }35 36        // Coordinate Compression37        Set<Integer> uniqueNumsSet = new HashSet<>();38        for (int num : nums) {39            uniqueNumsSet.add(num);40        }41        List<Integer> sortedUniqueNums = new ArrayList<>(uniqueNumsSet);42        Collections.sort(sortedUniqueNums);43        Map<Integer, Integer> valueToRank = new HashMap<>();44        for (int i = 0; i < sortedUniqueNums.size(); i++) {45            valueToRank.put(sortedUniqueNums.get(i), i + 1); // 1-based rank46        }47 48        int m = sortedUniqueNums.size();49        FenwickTree bit1 = new FenwickTree(m);50        FenwickTree bit2 = new FenwickTree(m);51        52        List<Integer> arr1 = new ArrayList<>();53        List<Integer> arr2 = new ArrayList<>();54 55        arr1.add(nums[0]);56        bit1.update(valueToRank.get(nums[0]), 1);57        58        arr2.add(nums[1]);59        bit2.update(valueToRank.get(nums[1]), 1);60 61        for (int i = 2; i < n; i++) {62            int val = nums[i];63            int rank = valueToRank.get(val);64 65            int count1 = arr1.size() - bit1.query(rank);66            int count2 = arr2.size() - bit2.query(rank);67 68            if (count1 > count2) {69                arr1.add(val);70                bit1.update(rank, 1);71            } else if (count1 < count2) {72                arr2.add(val);73                bit2.update(rank, 1);74            } else {75                if (arr1.size() <= arr2.size()) {76                    arr1.add(val);77                    bit1.update(rank, 1);78                } else {79                    arr2.add(val);80                    bit2.update(rank, 1);81                }82            }83        }84 85        int[] result = new int[n];86        int index = 0;87        for (int num : arr1) {88            result[index++] = num;89        }90        for (int num : arr2) {91            result[index++] = num;92        }93        return result;94    }95}

Complexity

Time

O(n log n). Coordinate compression takes `O(n log n)` due to sorting. The main loop runs `n` times, and each iteration involves BIT operations (`query` and `update`), which take `O(log m)` time, where `m <= n`. Thus, the loop takes `O(n log n)`. The total complexity is dominated by these parts.

Space

O(n). We store `arr1`, `arr2` (`O(n)`), data for coordinate compression (`O(m)`), and two Fenwick Trees (`O(m)`). Since `m <= n`, the total space is `O(n)`.

Trade-offs

Pros

Highly efficient, with a time complexity suitable for the given constraints.
A standard and powerful technique for problems involving rank-based queries.

Cons

More complex to implement than the brute-force approach.
Requires understanding of Fenwick Trees and coordinate compression.

Solutions

Solution

1class BinaryIndexedTree { private int n ; private int [] c ; public BinaryIndexedTree ( int n ) { this . n = n ; this . c = new int [ n + 1 ]; } public void update ( int x , int delta ) { for (; x <= n ; x += x & - x ) { c [ x ] += delta ; } } public int query ( int x ) { int s = 0 ; for (; x > 0 ; x -= x & - x ) { s += c [ x ]; } return s ; } } class Solution { public int [] resultArray ( int [] nums ) { int [] st = nums . clone (); Arrays . sort ( st ); int n = st . length ; BinaryIndexedTree tree1 = new BinaryIndexedTree ( n + 1 ); BinaryIndexedTree tree2 = new BinaryIndexedTree ( n + 1 ); tree1 . update ( Arrays . binarySearch ( st , nums [ 0 ]) + 1 , 1 ); tree2 . update ( Arrays . binarySearch ( st , nums [ 1 ]) + 1 , 1 ); int [] arr1 = new int [ n ]; int [] arr2 = new int [ n ]; arr1 [ 0 ] = nums [ 0 ]; arr2 [ 0 ] = nums [ 1 ]; int i = 1 , j = 1 ; for ( int k = 2 ; k < n ; ++ k ) { int x = Arrays . binarySearch ( st , nums [ k ]) + 1 ; int a = i - tree1 . query ( x ); int b = j - tree2 . query ( x ); if ( a > b ) { arr1 [ i ++] = nums [ k ]; tree1 . update ( x , 1 ); } else if ( a < b ) { arr2 [ j ++] = nums [ k ]; tree2 . update ( x , 1 ); } else if ( i <= j ) { arr1 [ i ++] = nums [ k ]; tree1 . update ( x , 1 ); } else { arr2 [ j ++] = nums [ k ]; tree2 . update ( x , 1 ); } } for ( int k = 0 ; k < j ; ++ k ) { arr1 [ i ++] = arr2 [ k ]; } return arr1 ; } }

1class BinaryIndexedTree { private: int n ; vector < int > c ; public: BinaryIndexedTree ( int n ) : n ( n ) , c ( n + 1 ) {} void update ( int x , int delta ) { for (; x <= n ; x += x & - x ) { c [ x ] += delta ; } } int query ( int x ) { int s = 0 ; for (; x > 0 ; x -= x & - x ) { s += c [ x ]; } return s ; } }; class Solution { public: vector < int > resultArray ( vector < int >& nums ) { vector < int > st = nums ; sort ( st . begin (), st . end ()); int n = st . size (); BinaryIndexedTree tree1 ( n + 1 ); BinaryIndexedTree tree2 ( n + 1 ); tree1 . update ( distance ( st . begin (), lower_bound ( st . begin (), st . end (), nums [ 0 ])) + 1 , 1 ); tree2 . update ( distance ( st . begin (), lower_bound ( st . begin (), st . end (), nums [ 1 ])) + 1 , 1 ); vector < int > arr1 = { nums [ 0 ]}; vector < int > arr2 = { nums [ 1 ]}; for ( int k = 2 ; k < n ; ++ k ) { int x = distance ( st . begin (), lower_bound ( st . begin (), st . end (), nums [ k ])) + 1 ; int a = arr1 . size () - tree1 . query ( x ); int b = arr2 . size () - tree2 . query ( x ); if ( a > b ) { arr1 . push_back ( nums [ k ]); tree1 . update ( x , 1 ); } else if ( a < b ) { arr2 . push_back ( nums [ k ]); tree2 . update ( x , 1 ); } else if ( arr1 . size () <= arr2 . size ()) { arr1 . push_back ( nums [ k ]); tree1 . update ( x , 1 ); } else { arr2 . push_back ( nums [ k ]); tree2 . update ( x , 1 ); } } arr1 . insert ( arr1 . end (), arr2 . begin (), arr2 . end ()); return arr1 ; } };

1class BinaryIndexedTree : __slots__ = "n" , "c" def __init__ ( self , n : int ): self . n = n self . c = [ 0 ] * ( n + 1 ) def update ( self , x : int , delta : int ) -> None : while x <= self . n : self . c [ x ] += delta x += x & - x def query ( self , x : int ) -> int : s = 0 while x : s += self . c [ x ] x -= x & - x return s class Solution : def resultArray ( self , nums : List [ int ]) -> List [ int ]: st = sorted ( set ( nums )) m = len ( st ) tree1 = BinaryIndexedTree ( m + 1 ) tree2 = BinaryIndexedTree ( m + 1 ) tree1 . update ( bisect_left ( st , nums [ 0 ]) + 1 , 1 ) tree2 . update ( bisect_left ( st , nums [ 1 ]) + 1 , 1 ) arr1 = [ nums [ 0 ]] arr2 = [ nums [ 1 ]] for x in nums [ 2 :]: i = bisect_left ( st , x ) + 1 a = len ( arr1 ) - tree1 . query ( i ) b = len ( arr2 ) - tree2 . query ( i ) if a > b : arr1 . append ( x ) tree1 . update ( i , 1 ) elif a < b : arr2 . append ( x ) tree2 . update ( i , 1 ) elif len ( arr1 ) <= len ( arr2 ): arr1 . append ( x ) tree1 . update ( i , 1 ) else : arr2 . append ( x ) tree2 . update ( i , 1 ) return arr1 + arr2

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem