Facebook Pixel

3072. Distribute Elements Into Two Arrays II

HardBinary Indexed TreeSegment TreeArraySimulation
Leetcode Link

Problem Description

You have a 1-indexed array nums of length n. You need to distribute all elements from this array into two new arrays arr1 and arr2 following specific rules.

The distribution process works as follows:

  • First operation: Put nums[1] into arr1
  • Second operation: Put nums[2] into arr2
  • For each subsequent element nums[i] (starting from index 3), you need to count how many elements in each array are strictly greater than nums[i]:
    • If arr1 has more elements greater than nums[i] than arr2, add nums[i] to arr1
    • If arr2 has more elements greater than nums[i] than arr1, add nums[i] to arr2
    • If both arrays have the same count of elements greater than nums[i]:
      • Add nums[i] to the array with fewer total elements
      • If both arrays have the same size, add nums[i] to arr1

After distributing all elements, concatenate arr1 and arr2 to form the result array. For example, if arr1 = [1,2,3] and arr2 = [4,5,6], then result = [1,2,3,4,5,6].

The solution uses Binary Indexed Trees (Fenwick Trees) to efficiently count elements greater than a given value. Since the value range can be large, the solution first discretizes the values by sorting unique elements and mapping them to smaller indices. Two Binary Indexed Trees track the elements in arr1 and arr2 respectively. For each new element, the solution queries how many elements are less than or equal to it, then calculates how many are greater by subtracting from the array length. This allows for efficient decision-making about which array should receive each element.

Quick Interview Experience
Help others by sharing your interview experience
Have you seen this problem before?

Intuition

The key challenge in this problem is efficiently counting how many elements in each array are greater than the current element we're trying to place. A naive approach would be to iterate through both arrays each time and count, but this would result in O(n²) time complexity.

We need a data structure that can:

  1. Handle insertions dynamically as we build the arrays
  2. Quickly answer range queries (how many elements are greater than a given value)

This naturally leads us to think about data structures that support efficient range queries. A Binary Indexed Tree (Fenwick Tree) is perfect for this scenario because it can update and query in O(log n) time.

However, there's another challenge: the values in the array could be very large (up to 10^9), and we can't create a Binary Indexed Tree of that size. This is where discretization comes in. Since we only care about the relative ordering of elements, not their actual values, we can map all unique values to a smaller range [1, m] where m is the number of unique values.

The insight for using Binary Indexed Trees specifically comes from recognizing that "count of elements greater than x" can be computed as: total_elements - count_of_elements_less_than_or_equal_to_x. Binary Indexed Trees excel at answering prefix sum queries (which gives us the count of elements less than or equal to a value when we treat each index as representing a discretized value).

By maintaining two separate Binary Indexed Trees (one for each array), we can track the distribution of elements in both arrays independently and make the placement decision in O(log n) time for each element. The discretization step ensures our trees are of manageable size while preserving the relative ordering needed for correct comparisons.

Learn more about Segment Tree patterns.

Solution Approach

The solution implements the strategy using discretization and Binary Indexed Trees. Let's walk through the implementation step by step:

1. Discretization Phase:

st = sorted(set(nums))
m = len(st)

First, we create a sorted list of unique values from nums. This maps the potentially large range of values to indices [0, m-1] where m is the number of unique values.

2. Initialize Data Structures:

tree1 = BinaryIndexedTree(m + 1)
tree2 = BinaryIndexedTree(m + 1)
arr1 = [nums[0]]
arr2 = [nums[1]]

We create two Binary Indexed Trees with size m + 1 (1-indexed) and initialize arr1 and arr2 with the first two elements as per the problem requirements.

3. Update Trees for Initial Elements:

tree1.update(bisect_left(st, nums[0]) + 1, 1)
tree2.update(bisect_left(st, nums[1]) + 1, 1)

We find the discretized positions of nums[0] and nums[1] using binary search (bisect_left) and update the respective trees. The +1 is because Binary Indexed Trees are 1-indexed.

4. Process Remaining Elements: For each element from nums[2] onwards:

i = bisect_left(st, x) + 1
a = len(arr1) - tree1.query(i)
b = len(arr2) - tree2.query(i)
  • Find the discretized position i of current element x
  • tree1.query(i) returns count of elements ≤ x in arr1
  • len(arr1) - tree1.query(i) gives count of elements > x in arr1
  • Similarly calculate for arr2

5. Decision Logic:

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)

Based on the greater count comparison:

  • If arr1 has more elements > x, add to arr1
  • If arr2 has more elements > x, add to arr2
  • If tied, add to the smaller array (or arr1 if equal sizes)

Binary Indexed Tree Operations:

  • update(x, delta): Adds delta to position x, propagating updates through the tree using x += x & -x
  • query(x): Returns sum of elements from positions 1 to x, traversing using x -= x & -x

Both operations work in O(log m) time due to the tree structure.

Time Complexity: O(n log n) for sorting + O(n log m) for processing = O(n log n)
Space Complexity: O(m) for the discretization array and Binary Indexed Trees

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example with nums = [5, 14, 3, 1, 2].

Step 1: Discretization

  • Extract unique values and sort: st = [1, 2, 3, 5, 14]
  • This maps our values to indices: 1→0, 2→1, 3→2, 5→3, 14→4

Step 2: Initialize Arrays and Trees

  • arr1 = [5] (first element)
  • arr2 = [14] (second element)
  • Update tree1 at position 4 (5 maps to index 3, +1 for 1-indexing)
  • Update tree2 at position 5 (14 maps to index 4, +1 for 1-indexing)

Step 3: Process nums[2] = 3

  • Find discretized position: 3 maps to index 2, so position = 3
  • Count elements > 3 in arr1: len(arr1) - tree1.query(3) = 1 - 0 = 1 (5 is greater)
  • Count elements > 3 in arr2: len(arr2) - tree2.query(3) = 1 - 0 = 1 (14 is greater)
  • Both have 1 element > 3, so check sizes: both arrays have size 1
  • Since sizes are equal, add to arr1: arr1 = [5, 3]
  • Update tree1 at position 3

Step 4: Process nums[3] = 1

  • Find discretized position: 1 maps to index 0, so position = 1
  • Count elements > 1 in arr1: len(arr1) - tree1.query(1) = 2 - 0 = 2 (both 5 and 3 are greater)
  • Count elements > 1 in arr2: len(arr2) - tree2.query(1) = 1 - 0 = 1 (14 is greater)
  • arr1 has more elements > 1, so add to arr1: arr1 = [5, 3, 1]
  • Update tree1 at position 1

Step 5: Process nums[4] = 2

  • Find discretized position: 2 maps to index 1, so position = 2
  • Count elements > 2 in arr1: len(arr1) - tree1.query(2) = 3 - 1 = 2 (5 and 3 are greater)
  • Count elements > 2 in arr2: len(arr2) - tree2.query(2) = 1 - 0 = 1 (14 is greater)
  • arr1 has more elements > 2, so add to arr1: arr1 = [5, 3, 1, 2]
  • Update tree1 at position 2

Final Result:

  • arr1 = [5, 3, 1, 2]
  • arr2 = [14]
  • Result = [5, 3, 1, 2, 14]

The Binary Indexed Trees efficiently tracked the distribution of values, allowing us to count elements greater than each new value in O(log m) time, where m is the number of unique values.

Solution Implementation

1from typing import List
2from bisect import bisect_left
3
4
5class BinaryIndexedTree:
6    """
7    Binary Indexed Tree (Fenwick Tree) for efficient prefix sum queries and updates.
8    Supports O(log n) update and query operations.
9    """
10    __slots__ = "n", "c"
11  
12    def __init__(self, n: int):
13        """
14        Initialize a Binary Indexed Tree with size n.
15      
16        Args:
17            n: The size of the tree (1-indexed)
18        """
19        self.n = n
20        self.c = [0] * (n + 1)  # Tree array (1-indexed)
21  
22    def update(self, x: int, delta: int) -> None:
23        """
24        Add delta to the value at position x.
25      
26        Args:
27            x: The position to update (1-indexed)
28            delta: The value to add
29        """
30        while x <= self.n:
31            self.c[x] += delta
32            x += x & -x  # Move to next node by adding lowest set bit
33  
34    def query(self, x: int) -> int:
35        """
36        Calculate the prefix sum from index 1 to x.
37      
38        Args:
39            x: The ending position for prefix sum (1-indexed)
40      
41        Returns:
42            The sum of elements from index 1 to x
43        """
44        prefix_sum = 0
45        while x > 0:
46            prefix_sum += self.c[x]
47            x -= x & -x  # Move to parent node by removing lowest set bit
48        return prefix_sum
49
50
51class Solution:
52    def resultArray(self, nums: List[int]) -> List[int]:
53        """
54        Divide the array into two subarrays based on counting elements greater than current.
55        Uses Binary Indexed Trees to efficiently count elements greater than a given value.
56      
57        Args:
58            nums: Input list of integers
59      
60        Returns:
61            Concatenation of the two resulting arrays
62        """
63        # Create sorted unique values for coordinate compression
64        sorted_unique = sorted(set(nums))
65        unique_count = len(sorted_unique)
66      
67        # Initialize two Binary Indexed Trees for the two arrays
68        tree1 = BinaryIndexedTree(unique_count + 1)
69        tree2 = BinaryIndexedTree(unique_count + 1)
70      
71        # Initialize first two elements
72        # Add 1 to indices since BIT is 1-indexed
73        first_idx = bisect_left(sorted_unique, nums[0]) + 1
74        second_idx = bisect_left(sorted_unique, nums[1]) + 1
75      
76        tree1.update(first_idx, 1)
77        tree2.update(second_idx, 1)
78      
79        arr1 = [nums[0]]
80        arr2 = [nums[1]]
81      
82        # Process remaining elements
83        for value in nums[2:]:
84            # Find compressed index for current value
85            compressed_idx = bisect_left(sorted_unique, value) + 1
86          
87            # Count elements strictly greater than current value in each array
88            # Elements greater = total elements - elements less than or equal
89            count_greater_arr1 = len(arr1) - tree1.query(compressed_idx)
90            count_greater_arr2 = len(arr2) - tree2.query(compressed_idx)
91          
92            # Decide which array to add the element to
93            if count_greater_arr1 > count_greater_arr2:
94                # More elements greater in arr1, add to arr1
95                arr1.append(value)
96                tree1.update(compressed_idx, 1)
97            elif count_greater_arr1 < count_greater_arr2:
98                # More elements greater in arr2, add to arr2
99                arr2.append(value)
100                tree2.update(compressed_idx, 1)
101            elif len(arr1) <= len(arr2):
102                # Equal count of greater elements, add to smaller array (arr1)
103                arr1.append(value)
104                tree1.update(compressed_idx, 1)
105            else:
106                # Equal count of greater elements, add to smaller array (arr2)
107                arr2.append(value)
108                tree2.update(compressed_idx, 1)
109      
110        # Return concatenation of both arrays
111        return arr1 + arr2
112
1/**
2 * Binary Indexed Tree (Fenwick Tree) for efficient range sum queries and point updates
3 */
4class BinaryIndexedTree {
5    private int size;
6    private int[] tree;
7
8    /**
9     * Initialize Binary Indexed Tree with given size
10     * @param size the maximum number of elements (1-indexed)
11     */
12    public BinaryIndexedTree(int size) {
13        this.size = size;
14        this.tree = new int[size + 1];
15    }
16
17    /**
18     * Update the value at position x by adding delta
19     * @param x the position to update (1-indexed)
20     * @param delta the value to add
21     */
22    public void update(int x, int delta) {
23        // Traverse all affected nodes in the tree
24        for (; x <= size; x += x & -x) {
25            tree[x] += delta;
26        }
27    }
28
29    /**
30     * Query the prefix sum from index 1 to x
31     * @param x the end position of the range (1-indexed)
32     * @return the sum of elements from 1 to x
33     */
34    public int query(int x) {
35        int sum = 0;
36        // Traverse the tree to calculate prefix sum
37        for (; x > 0; x -= x & -x) {
38            sum += tree[x];
39        }
40        return sum;
41    }
42}
43
44/**
45 * Solution class for array partitioning problem
46 */
47class Solution {
48    /**
49     * Partition the array into two subarrays based on the count of greater elements
50     * @param nums the input array
51     * @return the result array after partitioning and merging
52     */
53    public int[] resultArray(int[] nums) {
54        // Create a sorted copy for coordinate compression
55        int[] sortedNums = nums.clone();
56        Arrays.sort(sortedNums);
57        int n = sortedNums.length;
58      
59        // Initialize two Binary Indexed Trees for tracking elements in each array
60        BinaryIndexedTree tree1 = new BinaryIndexedTree(n + 1);
61        BinaryIndexedTree tree2 = new BinaryIndexedTree(n + 1);
62      
63        // Add first element to arr1 and second element to arr2
64        tree1.update(Arrays.binarySearch(sortedNums, nums[0]) + 1, 1);
65        tree2.update(Arrays.binarySearch(sortedNums, nums[1]) + 1, 1);
66      
67        // Initialize two arrays to store elements
68        int[] arr1 = new int[n];
69        int[] arr2 = new int[n];
70        arr1[0] = nums[0];
71        arr2[0] = nums[1];
72      
73        // Pointers for the next position in each array
74        int pointer1 = 1;
75        int pointer2 = 1;
76      
77        // Process remaining elements starting from index 2
78        for (int k = 2; k < n; ++k) {
79            // Get compressed coordinate (1-indexed)
80            int compressedValue = Arrays.binarySearch(sortedNums, nums[k]) + 1;
81          
82            // Count elements greater than current in each array
83            int greaterCount1 = pointer1 - tree1.query(compressedValue);
84            int greaterCount2 = pointer2 - tree2.query(compressedValue);
85          
86            // Decide which array to add the current element to
87            if (greaterCount1 > greaterCount2) {
88                // More greater elements in arr1, add to arr1
89                arr1[pointer1++] = nums[k];
90                tree1.update(compressedValue, 1);
91            } else if (greaterCount1 < greaterCount2) {
92                // More greater elements in arr2, add to arr2
93                arr2[pointer2++] = nums[k];
94                tree2.update(compressedValue, 1);
95            } else if (pointer1 <= pointer2) {
96                // Equal count, add to smaller or equal sized array (arr1)
97                arr1[pointer1++] = nums[k];
98                tree1.update(compressedValue, 1);
99            } else {
100                // Equal count and arr2 is smaller, add to arr2
101                arr2[pointer2++] = nums[k];
102                tree2.update(compressedValue, 1);
103            }
104        }
105      
106        // Append all elements from arr2 to arr1
107        for (int k = 0; k < pointer2; ++k) {
108            arr1[pointer1++] = arr2[k];
109        }
110      
111        return arr1;
112    }
113}
114
1class BinaryIndexedTree {
2private:
3    int size;
4    vector<int> tree;
5
6public:
7    // Initialize BIT with given size
8    BinaryIndexedTree(int size) 
9        : size(size), tree(size + 1) {}
10
11    // Add delta to element at position x (1-indexed)
12    void update(int x, int delta) {
13        // Traverse all affected nodes in BIT
14        for (; x <= size; x += x & -x) {
15            tree[x] += delta;
16        }
17    }
18
19    // Get prefix sum from index 1 to x (1-indexed)
20    int query(int x) {
21        int sum = 0;
22        // Traverse parents to calculate prefix sum
23        for (; x > 0; x -= x & -x) {
24            sum += tree[x];
25        }
26        return sum;
27    }
28};
29
30class Solution {
31public:
32    vector<int> resultArray(vector<int>& nums) {
33        // Create sorted copy for coordinate compression
34        vector<int> sortedNums = nums;
35        sort(sortedNums.begin(), sortedNums.end());
36        int n = sortedNums.size();
37      
38        // Initialize two BITs to track elements in each array
39        BinaryIndexedTree tree1(n + 1);
40        BinaryIndexedTree tree2(n + 1);
41      
42        // Find compressed coordinates for first two elements
43        int firstIdx = distance(sortedNums.begin(), 
44                               lower_bound(sortedNums.begin(), sortedNums.end(), nums[0])) + 1;
45        int secondIdx = distance(sortedNums.begin(), 
46                                lower_bound(sortedNums.begin(), sortedNums.end(), nums[1])) + 1;
47      
48        // Initialize first two elements
49        tree1.update(firstIdx, 1);
50        tree2.update(secondIdx, 1);
51      
52        vector<int> arr1 = {nums[0]};
53        vector<int> arr2 = {nums[1]};
54      
55        // Process remaining elements
56        for (int i = 2; i < n; ++i) {
57            // Get compressed coordinate for current element
58            int compressedIdx = distance(sortedNums.begin(), 
59                                        lower_bound(sortedNums.begin(), sortedNums.end(), nums[i])) + 1;
60          
61            // Count elements greater than current in each array
62            int greaterCount1 = arr1.size() - tree1.query(compressedIdx);
63            int greaterCount2 = arr2.size() - tree2.query(compressedIdx);
64          
65            // Decide which array to add current element to
66            if (greaterCount1 > greaterCount2) {
67                // More greater elements in arr1, add to arr1
68                arr1.push_back(nums[i]);
69                tree1.update(compressedIdx, 1);
70            } else if (greaterCount1 < greaterCount2) {
71                // More greater elements in arr2, add to arr2
72                arr2.push_back(nums[i]);
73                tree2.update(compressedIdx, 1);
74            } else if (arr1.size() <= arr2.size()) {
75                // Equal greater elements, add to smaller or equal sized array
76                arr1.push_back(nums[i]);
77                tree1.update(compressedIdx, 1);
78            } else {
79                // arr2 is smaller, add to arr2
80                arr2.push_back(nums[i]);
81                tree2.update(compressedIdx, 1);
82            }
83        }
84      
85        // Merge arr2 into arr1 and return result
86        arr1.insert(arr1.end(), arr2.begin(), arr2.end());
87        return arr1;
88    }
89};
90
1// Binary Indexed Tree (Fenwick Tree) for efficient range sum queries
2// Array to store the tree values
3let treeSize: number;
4let treeArray: number[];
5
6// Initialize the Binary Indexed Tree with given size
7function initializeBIT(size: number): void {
8    treeSize = size;
9    treeArray = Array(size + 1).fill(0);
10}
11
12// Update the tree by adding delta to position x
13// Uses the property that x & -x gives the rightmost set bit
14function updateBIT(position: number, delta: number): void {
15    while (position <= treeSize) {
16        treeArray[position] += delta;
17        position += position & -position; // Move to next position to update
18    }
19}
20
21// Query the cumulative sum from index 1 to x
22function queryBIT(position: number): number {
23    let sum = 0;
24    while (position > 0) {
25        sum += treeArray[position];
26        position -= position & -position; // Move to parent node
27    }
28    return sum;
29}
30
31// Binary search to find the leftmost position where sortedArray[mid] >= target
32function binarySearch(sortedArray: number[], target: number): number {
33    let left = 0;
34    let right = sortedArray.length;
35  
36    while (left < right) {
37        const mid = Math.floor((left + right) / 2);
38        if (sortedArray[mid] >= target) {
39            right = mid;
40        } else {
41            left = mid + 1;
42        }
43    }
44    return left;
45}
46
47// Main function to split array based on greater element count
48function resultArray(nums: number[]): number[] {
49    // Create a sorted copy of nums for coordinate compression
50    const sortedNums: number[] = [...nums].sort((a, b) => a - b);
51    const arrayLength: number = sortedNums.length;
52  
53    // Initialize two Binary Indexed Trees for tracking elements in each array
54    const tree1Size = arrayLength + 1;
55    const tree2Size = arrayLength + 1;
56  
57    // Tree 1 for first array
58    let tree1Array: number[] = Array(tree1Size + 1).fill(0);
59    let currentTree1Size = tree1Size;
60  
61    // Tree 2 for second array
62    let tree2Array: number[] = Array(tree2Size + 1).fill(0);
63    let currentTree2Size = tree2Size;
64  
65    // Helper functions for tree1
66    const updateTree1 = (position: number, delta: number): void => {
67        while (position <= currentTree1Size) {
68            tree1Array[position] += delta;
69            position += position & -position;
70        }
71    };
72  
73    const queryTree1 = (position: number): number => {
74        let sum = 0;
75        while (position > 0) {
76            sum += tree1Array[position];
77            position -= position & -position;
78        }
79        return sum;
80    };
81  
82    // Helper functions for tree2
83    const updateTree2 = (position: number, delta: number): void => {
84        while (position <= currentTree2Size) {
85            tree2Array[position] += delta;
86            position += position & -position;
87        }
88    };
89  
90    const queryTree2 = (position: number): number => {
91        let sum = 0;
92        while (position > 0) {
93            sum += tree2Array[position];
94            position -= position & -position;
95        }
96        return sum;
97    };
98  
99    // Initialize the two result arrays with first two elements
100    const firstArray: number[] = [nums[0]];
101    const secondArray: number[] = [nums[1]];
102  
103    // Add first element to tree1 (1-indexed)
104    const firstElementIndex = binarySearch(sortedNums, nums[0]) + 1;
105    updateTree1(firstElementIndex, 1);
106  
107    // Add second element to tree2 (1-indexed)
108    const secondElementIndex = binarySearch(sortedNums, nums[1]) + 1;
109    updateTree2(secondElementIndex, 1);
110  
111    // Process remaining elements starting from index 2
112    for (let idx = 2; idx < nums.length; idx++) {
113        const currentValue = nums[idx];
114        // Get compressed coordinate (1-indexed)
115        const compressedIndex = binarySearch(sortedNums, currentValue) + 1;
116      
117        // Count elements greater than current value in each array
118        // Total elements - elements less than or equal to current
119        const greaterCountInFirst = firstArray.length - queryTree1(compressedIndex);
120        const greaterCountInSecond = secondArray.length - queryTree2(compressedIndex);
121      
122        // Decide which array to add the current element to
123        if (greaterCountInFirst > greaterCountInSecond) {
124            // More greater elements in first array, add to first
125            firstArray.push(currentValue);
126            updateTree1(compressedIndex, 1);
127        } else if (greaterCountInFirst < greaterCountInSecond) {
128            // More greater elements in second array, add to second
129            secondArray.push(currentValue);
130            updateTree2(compressedIndex, 1);
131        } else {
132            // Equal count of greater elements, add to smaller array
133            if (firstArray.length <= secondArray.length) {
134                firstArray.push(currentValue);
135                updateTree1(compressedIndex, 1);
136            } else {
137                secondArray.push(currentValue);
138                updateTree2(compressedIndex, 1);
139            }
140        }
141    }
142  
143    // Concatenate both arrays and return result
144    return firstArray.concat(secondArray);
145}
146

Time and Space Complexity

Time Complexity: O(n × log n)

The time complexity analysis breaks down as follows:

  • Creating the sorted set st = sorted(set(nums)) takes O(n log n) time
  • Initializing the two Binary Indexed Trees takes O(m) time where m ≤ n
  • Processing the first two elements involves:
    • bisect_left operations: O(log m) each
    • BIT update operations: O(log m) each
  • The main loop iterates through n-2 elements, and for each element:
    • bisect_left(st, x): O(log m) time
    • tree1.query(i) and tree2.query(i): O(log m) time each
    • One BIT update operation: O(log m) time
  • Since the loop runs O(n) times with O(log m) operations per iteration, and m ≤ n, the loop contributes O(n log n)
  • The final concatenation arr1 + arr2 takes O(n) time

Overall: O(n log n) + O(n × log n) + O(n) = O(n × log n)

Space Complexity: O(n)

The space complexity analysis:

  • The sorted set st stores at most n unique elements: O(n)
  • Two Binary Indexed Trees, each with size m+1 ≤ n+1: O(n)
  • Arrays arr1 and arr2 together store all n elements: O(n)
  • Additional variables use constant space: O(1)

Total space complexity: O(n) + O(n) + O(n) = O(n)

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

Common Pitfalls

1. Off-by-One Errors with 1-Indexed Binary Indexed Trees

One of the most common mistakes when implementing this solution is forgetting that Binary Indexed Trees are traditionally 1-indexed, while Python lists are 0-indexed. This can lead to incorrect index calculations when updating or querying the tree.

Pitfall Example:

# WRONG: Forgetting to add 1 when converting to BIT index
compressed_idx = bisect_left(sorted_unique, value)  # 0-indexed
tree1.update(compressed_idx, 1)  # BIT expects 1-indexed!

Solution: Always add 1 when converting from the 0-indexed result of bisect_left to the 1-indexed BIT position:

# CORRECT: Convert 0-indexed to 1-indexed
compressed_idx = bisect_left(sorted_unique, value) + 1
tree1.update(compressed_idx, 1)

2. Incorrectly Calculating Count of Elements Greater Than Current

Another frequent error is miscalculating the number of elements strictly greater than the current value. The BIT's query(i) returns the count of elements less than or equal to the value at index i, not strictly less than.

Pitfall Example:

# WRONG: This counts elements >= value, not > value
count_greater = len(arr1) - tree1.query(compressed_idx - 1)

Solution: Use the correct formula: total elements minus elements less than or equal to current:

# CORRECT: Count elements strictly greater than value
count_greater = len(arr1) - tree1.query(compressed_idx)

3. Forgetting to Handle Duplicate Values in Discretization

When discretizing values, using sorted(nums) instead of sorted(set(nums)) creates duplicate entries, leading to incorrect index mappings and wasted space.

Pitfall Example:

# WRONG: Duplicates in sorted array cause incorrect mappings
sorted_values = sorted(nums)  # [1, 1, 2, 2, 3] - duplicates!
idx = bisect_left(sorted_values, 2)  # May return different indices for same value

Solution: Always use set() to remove duplicates before sorting:

# CORRECT: Remove duplicates for proper discretization
sorted_unique = sorted(set(nums))  # [1, 2, 3] - no duplicates
idx = bisect_left(sorted_unique, 2)  # Consistent index for value 2

4. Incorrect Tie-Breaking Logic

When the count of greater elements is equal in both arrays, the tie-breaking rule must check array sizes correctly. A common mistake is using the wrong comparison operator or forgetting this step entirely.

Pitfall Example:

# WRONG: Using wrong comparison or always defaulting to arr1
if count_greater_arr1 == count_greater_arr2:
    arr1.append(value)  # Always adding to arr1 regardless of size!
    tree1.update(compressed_idx, 1)

Solution: Implement proper tie-breaking by comparing array lengths:

# CORRECT: Check array sizes when counts are equal
elif len(arr1) <= len(arr2):
    arr1.append(value)
    tree1.update(compressed_idx, 1)
else:
    arr2.append(value)
    tree2.update(compressed_idx, 1)
Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Recommended Readings

Segment Tree

Segment Tree Faster Range Queries For this article we want to introduce the idea of a Segment Tree Segment Trees allow us to quickly perform range queries as well as range updates Suppose we have an array and we want to know the sum of a particular range of numbers as well

Coding Interview Patterns: The Smart, Data-Driven Way to Prepare for Tech Interviews

Coding Interview Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way

Recursion Review

Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https assets algo monster recursion jpg You first call Ben and ask

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More