3072. Distribute Elements Into Two Arrays II

HardBinary Indexed TreeSegment TreeArraySimulation
Leetcode Link

Problem Description

You are given an array of integers named nums, and it is 1-indexed, meaning the indexing starts from 1 instead of 0. The length of the array is n. There's a special function called greaterCount which, given an array arr and a value val, returns how many elements in arr are strictly larger than val.

The task is to distribute all the elements of the array nums into two new arrays arr1 and arr2. This should be done over n operations aligning with the following rules:

  1. For the first operation, put nums[1] into arr1.
  2. For the second operation, put nums[2] into arr2.
  3. From the third operation onward, decide whether to put nums[i] into arr1 or arr2 based on comparing greaterCount(arr1, nums[i]) to greaterCount(arr2, nums[i]):
    • If arr1 has more elements greater than nums[i] than arr2 does, put nums[i] into arr1.
    • If arr2 has more elements greater than nums[i] than arr1 does, put nums[i] into arr2.
    • If both arrays have an equal number of elements greater than nums[i], put nums[i] into the one with fewer elements.
    • If both arrays are equal in size, put nums[i] into arr1.

Finally, the output should be an array result consisting of all the elements first from arr1 and then from arr2.

Intuition

The problem adds some complexity because you are not just splitting the array, but you're also required to compare elements in a specific way during the process. A Naive approach, like comparing each element with every other element to calculate the greaterCount every time, would be too slow, especially for large arrays.

To efficiently solve this problem, we can use a Binary Indexed Tree (BIT), also known as a Fenwick Tree. This data structure allows us to update elements and calculate prefix sums in logarithmic time, which is much faster than the linear time a naive approach would require.

However, there's a catch. The numbers in nums might be too large or sparse to use directly in a BIT, which requires index-based access. This is where discretization comes into play. Discretization involves mapping the elements to a compact range of indices such that relative order among the elements is preserved. This can be done by sorting the unique elements of nums and using their index in the sorted order as a new index.

By discretizing the numbers, we can use the indices in the BIT without having to deal with large numbers directly. Now, the greaterCount can be found by subtracting the query result (which gives us the count of numbers less than or equal to the current number) from the total number of elements already in the array. Then we can compare the counts to decide where to put the current number. By keeping a BIT for both arr1 and arr2, we can maintain the necessary information to perform each operation efficiently as described.

The BITs (for both arr1 and arr2) are utilized to maintain a frequency count of the elements as they are added. This allows us to track the criteria specified in the operations and place elements according to the greaterCount comparison, or by size if counts are equal. With this approach, we can construct the correct split of nums into arr1 and arr2.

Learn more about Segment Tree patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Which algorithm is best for finding the shortest distance between two points in an unweighted graph?

Solution Approach

The solution uses a combination of discretization and two Binary Indexed Trees (BITs) to efficiently distribute the elements of the nums array into two arrays, arr1 and arr2.

Discretization

Discretization is used to handle the potentially large or sparse integers in nums. The process consists of the following steps:

  • Create a sorted set of the unique elements in nums. This removes duplicates and orders the elements.
  • Use the index of each element in this sorted set as its new "discrete" index. This maps the wide range of nums into a compact range suitable for index-based data structures, such as the BIT.

Binary Indexed Tree

The Binary Indexed Tree, or BIT, is a data structure that allows us to:

  • Update the frequency count of elements (update operation) efficiently.
  • Query the cumulative frequency up to a certain index (query operation) efficiently.

These operations both run in O(log n) time, which is significantly faster than a naive approach that might have O(n) complexity.

Implementation Details

  1. Two instances of the BIT are created, tree1 for arr1 and tree2 for arr2, using the length of the discretized array plus one to accommodate one-based indexing.

  2. The algorithm begins by placing the first element of nums into arr1 and the second element into arr2. It then updates the corresponding BITs by increasing the frequency count at the index of each element in the discretized set. This is done using the update function.

  3. Starting from the third element, the algorithm performs the following steps for each nums[i]:

    • Discretize nums[i] to find its index in the sorted set.
    • Query both BITs (tree1 and tree2) to get the number of elements less than or equal to nums[i] in both arrays. The result of this query is subtracted from the length of the respective array to get the greaterCount.
    • Compare greaterCount(arr1, nums[i]) to greaterCount(arr2, nums[i]), and based on this comparison and the lengths of arr1 and arr2, decide where to append the current element.
    • Update the corresponding BIT to reflect the addition of the new element.
  4. Continue this process until all the elements from nums have been distributed between arr1 and arr2.

  5. Concatenate arr1 and arr2 to form the result array, which is then returned.

By using this approach, each operation enforces the rules for distributing elements between arr1 and arr2 as stated in the problem description, ensuring the correct construction of the result array.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

What is the worst case running time for finding an element in a binary tree (not necessarily binary search tree) of size n?

Example Walkthrough

Let's walk through an example using the solution approach described. Suppose we have the following nums array:

1nums = [4, 7, 3, 1, 9, 2, 5]

Step 1: Discretization

  1. Create a sorted set of the unique elements in nums:

    1sorted_unique_nums = [1, 2, 3, 4, 5, 7, 9]
  2. Map each element to its index in the sorted array:

    1discretized_map = {
    2    4: 4,
    3    7: 6,
    4    3: 3,
    5    1: 1,
    6    9: 7,
    7    2: 2,
    8    5: 5
    9}

Step 2: Initialize the Binary Indexed Trees

Initialize tree1 and tree2 with 8 nodes (to accommodate 7 unique discretized indices plus 1 for one-based indexing).

Step 3: Begin Distribution

  • For the first operation, place nums[1] (which is 4) into arr1 and update tree1:

    1arr1 = [4]
    2arr2 = []

    The discretized value of 4 is 4, so we update tree1[4].

  • For the second operation, place nums[2] (which is 7) into arr2 and update tree2:

    1arr1 = [4]
    2arr2 = [7]

    The discretized value of 7 is 6, so we update tree2[6].

Step 4: Perform Subsequent Operations

  • For the third element nums[3] (which is 3), we do the following:

    1. Discretize 3 to get 3.
    2. Query tree1 and tree2 to get counts of elements less than or equal to 3. Both will be 0 because there are no elements less than or equal 3 added yet.
    3. Since no elements are greater in either array, compare the lengths of arr1 and arr2. Both have the same length, so place 3 into arr1 as per the rules.
    4. Update tree1 by increasing the frequency at index 3 in the tree.

    Now arr1 and arr2 looks like this:

    1arr1 = [4, 3]
    2arr2 = [7]

Continuing in the same fashion:

  • For nums[4] (1):

    1arr1 = [4, 3]
    2arr2 = [7, 1]
  • For nums[5] (9):

    1arr1 = [4, 3]
    2arr2 = [7, 1, 9]
  • For nums[6] (2):

    1arr1 = [4, 3, 2]
    2arr2 = [7, 1, 9]
  • Finally for nums[7] (5):

    1arr1 = [4, 3, 2, 5]
    2arr2 = [7, 1, 9]

Step 5: Combine arr1 and arr2 to form the result array

1result = arr1 + arr2 = [4, 3, 2, 5, 7, 1, 9]

Conclusion

Using discretization and the Binary Indexed Trees, we have efficiently distributed the elements into arr1 and arr2 while following the rules given in the problem description. The final output is the result array.

Solution Implementation

1from typing import List
2from bisect import bisect_left
3
4class BinaryIndexedTree:
5    # Using __slots__ for memory optimization since we are sure 
6    # about what attributes the class instances will hold
7    __slots__ = "size", "tree"
8
9    def __init__(self, n: int):
10        """
11        Initialize the Binary Indexed Tree with a specified size.
12      
13        :param n: The size of the array for which the tree is constructed.
14        """
15        self.size = n
16        self.tree = [0] * (n + 1)
17
18    def update(self, index: int, delta: int) -> None:
19        """
20        Update the Binary Indexed Tree at the given index by adding the delta.
21      
22        :param index: The index in the array to update.
23        :param delta: The value to add to the Binary Indexed Tree at the index.
24        """
25        while index <= self.size:
26            self.tree[index] += delta
27            index += index & -index
28
29    def query(self, index: int) -> int:
30        """
31        Calculate the prefix sum up to the given index in the Binary Indexed Tree.
32      
33        :param index: The index up to which the prefix sum is computed.
34        :return: The prefix sum.
35        """
36        sum_ = 0
37        while index:
38            sum_ += self.tree[index]
39            index -= index & -index
40        return sum_
41
42
43class Solution:
44    def resultArray(self, nums: List[int]) -> List[int]:
45        """
46        Divide a list of numbers into two sequences based on counts
47        of numbers less than each number using Binary Indexed Trees.
48      
49        :param nums: List of integers to be divided.
50        :return: The combined result from both sequences.
51        """
52        # Sort and remove duplicates to map values to tree indexes
53        sorted_unique_nums = sorted(set(nums))
54        mapped_length = len(sorted_unique_nums)
55      
56        # Create two Binary Indexed Trees
57        tree1 = BinaryIndexedTree(mapped_length)
58        tree2 = BinaryIndexedTree(mapped_length)
59      
60        # Initialize the sequences with the first two numbers
61        index1 = bisect_left(sorted_unique_nums, nums[0]) + 1
62        index2 = bisect_left(sorted_unique_nums, nums[1]) + 1
63        tree1.update(index1, 1)
64        tree2.update(index2, 1)
65        seq1 = [nums[0]]
66        seq2 = [nums[1]]
67
68        # Process the rest of the numbers
69        for x in nums[2:]:
70            index = bisect_left(sorted_unique_nums, x) + 1
71            count1 = len(seq1) - tree1.query(index)
72            count2 = len(seq2) - tree2.query(index)
73          
74            # Decide to which sequence to add the current number based on the counts
75            if count1 > count2 or (count1 == count2 and len(seq1) <= len(seq2)):
76                seq1.append(x)
77                tree1.update(index, 1)
78            else:
79                seq2.append(x)
80                tree2.update(index, 1)
81      
82        # Combine both sequences
83        return seq1 + seq2
84
1import java.util.Arrays;
2
3// Class representing a Binary Indexed Tree (Fenwick Tree)
4class BinaryIndexedTree {
5    private int size; // the number of elements in the Binary Indexed Tree
6    private int[] tree; // array that represents the Binary Indexed Tree
7
8    // Constructor that initializes the tree with a given size
9    public BinaryIndexedTree(int size) {
10        this.size = size;
11        this.tree = new int[size + 1];
12    }
13
14    // Updates a value at the given index by a certain delta amount
15    public void update(int index, int delta) {
16        // Loop over the tree and apply updates
17        for (; index <= size; index += index & -index) {
18            tree[index] += delta;
19        }
20    }
21
22    // Queries and returns the cumulative frequency up to a given index
23    public int query(int index) {
24        int sum = 0;
25        // Loop over the tree and calculate the sum
26        for (; index > 0; index -= index & -index) {
27            sum += tree[index];
28        }
29        return sum;
30    }
31}
32
33// Solution class that contains the method for creating the result array
34class Solution {
35    // Method to generate the result array based on certain constraints
36    public int[] resultArray(int[] nums) {
37        // Create a sorted copy of the original array
38        int[] sortedNums = nums.clone();
39        Arrays.sort(sortedNums);
40        int n = sortedNums.length;
41      
42        // Set up two Binary Indexed Trees for processing the arrays
43        BinaryIndexedTree tree1 = new BinaryIndexedTree(n + 1);
44        BinaryIndexedTree tree2 = new BinaryIndexedTree(n + 1);
45      
46        // Initial updates for the first two elements of the nums array
47        tree1.update(Arrays.binarySearch(sortedNums, nums[0]) + 1, 1);
48        tree2.update(Arrays.binarySearch(sortedNums, nums[1]) + 1, 1);
49      
50        // Initialize result arrays and their pointers
51        int[] resultArray1 = new int[n];
52        int[] resultArray2 = new int[n];
53      
54        // Set the first elements of the result arrays
55        resultArray1[0] = nums[0];
56        resultArray2[0] = nums[1];
57      
58        // Pointers for the result arrays
59        int index1 = 1; 
60        int index2 = 1;
61      
62        // Process the rest of the elements in the nums array
63        for (int k = 2; k < n; ++k) {
64            int valuePosition = Arrays.binarySearch(sortedNums, nums[k]) + 1;
65            int a = index1 - tree1.query(valuePosition);
66            int b = index2 - tree2.query(valuePosition);
67            // Compare and decide which array to add the current element to
68            if (a > b) {
69                resultArray1[index1++] = nums[k];
70                tree1.update(valuePosition, 1);
71            } else if (a < b) {
72                resultArray2[index2++] = nums[k];
73                tree2.update(valuePosition, 1);
74            } else if (index1 <= index2) {
75                resultArray1[index1++] = nums[k];
76                tree1.update(valuePosition, 1);
77            } else {
78                resultArray2[index2++] = nums[k];
79                tree2.update(valuePosition, 1);
80            }
81        }
82
83        // Merge the two result arrays into resultArray1
84        for (int k = 0; k < index2; ++k) {
85            resultArray1[index1++] = resultArray2[k];
86        }
87      
88        // Return the merged result array
89        return resultArray1;
90    }
91}
92
1#include <vector>
2#include <algorithm>
3using namespace std;
4
5// BinaryIndexedTree (also known as a Fenwick tree) for efficient 
6// update and query of prefix sums.
7class BinaryIndexedTree {
8private:
9    int size;  // Total number of elements in the array
10    vector<int> tree;  // The tree structure stored as a vector
11
12public:
13    // Constructor to initialize the tree with a given size.
14    BinaryIndexedTree(int size)
15        : size(size), tree(size + 1, 0) {}
16
17    // Updates the tree with a given delta at a specific index.
18    void update(int index, int delta) {
19        for (; index <= size; index += index & -index) {
20            tree[index] += delta;
21        }
22    }
23
24    // Queries the prefix sum up to a given index.
25    int query(int index) {
26        int sum = 0;
27        for (; index > 0; index -= index & -index) {
28            sum += tree[index];
29        }
30        return sum;
31    }
32};
33
34// Solution class containing the method to process the input array
35// and produce the result array.
36class Solution {
37public:
38    // Method to create and return the result array.
39    vector<int> resultArray(vector<int>& nums) {
40        // Copy the input array and sort the copy to facilitate binary index calculations.
41        vector<int> sortedNums = nums;
42        sort(sortedNums.begin(), sortedNums.end());
43
44        int n = nums.size();  // Store the size of the input array.
45      
46        // Create two BinaryIndexedTree instances.
47        BinaryIndexedTree tree1(n + 1);
48        BinaryIndexedTree tree2(n + 1);
49
50        // Update the trees with the first two elements in nums.
51        tree1.update(distance(sortedNums.begin(), lower_bound(sortedNums.begin(), sortedNums.end(), nums[0])) + 1, 1);
52        tree2.update(distance(sortedNums.begin(), lower_bound(sortedNums.begin(), sortedNums.end(), nums[1])) + 1, 1);
53
54        // Initialize two arrays to store elements as we process them.
55        vector<int> arr1 = {nums[0]};
56        vector<int> arr2 = {nums[1]};
57
58        // Process the remaining elements in nums.
59        for (int k = 2; k < n; ++k) {
60            int x = distance(sortedNums.begin(), lower_bound(sortedNums.begin(), sortedNums.end(), nums[k])) + 1;
61            int a = arr1.size() - tree1.query(x);
62            int b = arr2.size() - tree2.query(x);
63            if (a > b) {
64                arr1.push_back(nums[k]);
65                tree1.update(x, 1);
66            } else if (a < b) {
67                arr2.push_back(nums[k]);
68                tree2.update(x, 1);
69            } else if (arr1.size() <= arr2.size()) {
70                arr1.push_back(nums[k]);
71                tree1.update(x, 1);
72            } else {
73                arr2.push_back(nums[k]);
74                tree2.update(x, 1);
75            }
76        }
77
78        // Combine arr1 and arr2 into arr1 as the final result.
79        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
80        return arr1;
81    }
82};
83
1const n: number = 0; // Size of the Binary Indexed Tree
2const bitValues: number[] = []; // Array that represents the tree
3
4// Initializes the Binary Indexed Tree with a specified size
5function initializeBIT(size: number): void {
6    this.n = size;
7    this.bitValues = Array(size + 1).fill(0);
8}
9
10// Updates the Binary Indexed Tree at a specific index by a delta value
11function updateBIT(index: number, delta: number): void {
12    while (index <= n) {
13        bitValues[index] += delta;
14        index += (index & -index);
15    }
16}
17
18// Queries the sum of the range from 1 to a specific index in the Binary Indexed Tree
19function queryBIT(index: number): number {
20    let sum = 0;
21    while (index > 0) {
22        sum += bitValues[index];
23        index -= (index & -index);
24    }
25    return sum;
26}
27
28// Given an array of numbers, returns the resulting array after operations are performed
29function getResultArray(nums: number[]): number[] {
30    const numSorted: number[] = nums.slice().sort((a, b) => a - b);
31    const search = (value: number): number => {
32        let left = 0, right = numSorted.length;
33        while (left < right) {
34            const mid = (left + right) >> 1;
35            if (numSorted[mid] >= value) {
36                right = mid;
37            } else {
38                left = mid + 1;
39            }
40        }
41        return left;
42    };
43    initializeBIT(numSorted.length + 1);
44    const bit1: number[] = Array.from(bitValues);
45    const bit2: number[] = Array.from(bitValues);
46    updateBIT(search(nums[0]) + 1, 1);
47    updateBIT(search(nums[1]) + 1, 1);
48    const arr1: number[] = [nums[0]];
49    const arr2: number[] = [nums[1]];
50    nums.slice(2).forEach(x => {
51        const index: number = search(x) + 1;
52        const countA: number = arr1.length - queryBIT(index); // using global queryBIT
53        const countB: number = arr2.length - queryBIT(index); // using global queryBIT
54        if (countA > countB || (countA === countB && arr1.length <= arr2.length)) {
55            arr1.push(x);
56            updateBIT(index, 1); // using global updateBIT
57        } else {
58            arr2.push(x);
59            updateBIT(index, 1); // using global updateBIT
60        }
61    });
62    return arr1.concat(arr2);
63}
64
65// Note that the code uses this.n and this.bitValues assuming they are members of a BIT class,
66// but the problem statement asks to omit the class, hence global variables n and bitValues are
67// introduced and would need to be passed appropriately if this code is part of a larger system.
68
Not Sure What to Study? Take the 2-min Quiz:

What's the relationship between a tree and a graph?

Time and Space Complexity

The time complexity of the BinaryIndexedTree operations update and query is O(log n), where n is the length of the nums array. Inside the function resultArray, update is called once per element, and query is also called once per element when constructing the arr1 and arr2. This leads to a time complexity of O(n log n) because there are n elements processed and for each element, an O(log n) operation is performed.

The sorted function and the conversion of nums to a set has a time complexity of O(n log n) because it sorts the unique elements in the nums list.

The bisect_left function has a time complexity of O(log n) since it performs binary search on the sorted unique elements. It is called within the for loop, so it does not dominate the time complexity of the algorithm, which remains O(n log n).

The space complexity is O(n) as the BinaryIndexedTree requires additional space proportional to the number of unique elements, and the temporary arrays arr1 and arr2 store elements from nums. Since all elements of nums could potentially be unique, the space required may scale linearly with n.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

What is the running time of the following code?

1int sqrt(int n) {
2  for (int guess = 1; guess * guess <= n; guess++) {
3    if (guess * guess == n) {
4      return guess;
5    }
6  }
7  return -1;
8}

Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫