2602. Minimum Operations to Make All Array Elements Equal

MediumArrayBinary SearchPrefix SumSorting

Problem Description

In this problem, we are given an array of positive integers called nums and a separate array of positive integers called queries. For each query queries[i], the goal is to transform all elements in nums such that they are all equal to the value of queries[i]. To do this, we can increase or decrease any element in the nums array by 1. The objective is to find the minimum number of these increase or decrease operations required to achieve the goal for each query.

It is important to mention that these operations are independent for each query; that is, after performing the operations required by one query, the nums array returns to its original state before applying the next query.

Our task is to return an array answer where each element answer[i] is the minimum number of operations needed to make all elements of nums equal to queries[i].

Intuition

The straightforward approach of repeatedly increasing or decreasing each number in nums until it matches queries[i] would be too slow, especially for large arrays.

To optimize this, we can think about the problem in two parts:

We need to reduce the larger numbers down to queries[i].
We need to increase the smaller numbers up to queries[i].

Each of these parts needs a different sum of operations. If the numbers are sorted, we can quickly find the point in the array where the numbers shift from being less than to greater than queries[i]. This helps us to separate the two groups conveniently.

Once we have the numbers sorted, we can use a prefix sum array (s in the solution code) to quickly calculate the total sum of numbers up to any point. The prefix sum array keeps a running total, which allows us to compute the sum of any sub-array in constant time, by subtracting two prefix sums.

With the help of binary search, we can find the exact point of separation between numbers we need to increase and those we need to decrease. The binary search efficiently finds the position in a sorted array where a given element could be inserted to maintain the order (less than or equal / greater than queries[i] in our case).

Using the point of separation and prefix sums, we can calculate the number of operations needed for both increasing the smaller numbers and decreasing the larger ones. The sum of these two gives us the minimum number of operations needed for the nums array to be all equal to queries[i].

Finally, we repeat this process for each query and compile the results into an answer array.

Learn more about Binary Search, Prefix Sum and Sorting patterns.

Solution Approach

The solution employs a combination of sorting, prefix summation, and binary search to efficiently find the minimum number of operations required to make all elements of nums equal to each query value.

Here is a detailed breakdown of the solution approach:

Sorting nums: We start by sorting the array nums. Sorting is helpful because it allows us to apply binary search later in the process and also makes it easy to separate elements into those that need to be increased vs. those that need to be decreased to match the query value.
Prefix Sum Array (s): After sorting nums, we compute the prefix sum array s. The prefix sum is essentially a cumulative sum up to the current index, and we initialize it with a zero to account for the hypothetical prefix sum before the first element. The formula used to obtain the prefix sum array is s[i] = s[i-1] + nums[i-1]. This approach simplifies calculating the sum of subarrays later on.
Binary Search: For each query x, we need to find the split point where we increase the numbers below x and decrease the numbers above x. We use the bisect_left function from the bisect module, which is an implementation of the binary search algorithm, to find this index quickly. We find two indices in this step - the index of the first element greater than x and the index of the first element equal to or greater than x.
Calculating the Number of Operations:
- For the elements greater than x, we subtract the sum of these elements given by s[-1] - s[i] from the total reduction required (len(nums) - i) * x to find the total decrease operations needed t.
- For elements less than x, we multiply x * i to find the total increase needed and subtract the sum of these elements s[i] to find the total increase operations needed.
- The sum of these two values gives us the total minimum operations required for current query x.
Build the Result Array (ans): We append the total minimum operations found for each query to the result array ans.

Combining these steps and repeating them for each value in queries, we arrive at the final result of the minimum number of operations required to equalize the nums array to each of the query values. The use of binary search and prefix sums ensures that the operations are performed efficiently, managing to keep the computational complexity within acceptable limits even for large arrays.

Ready to land your dream job?

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

Start Evaluator

Example Walkthrough

Let's illustrate the solution approach using an example. Given an array nums = [1, 3, 4, 9] and a queries array containing a single element [6], follow the steps below to find the minimum number of operations required to make all elements in nums equal to 6.

Sorting nums: The array nums is already sorted: [1, 3, 4, 9].
Prefix Sum Array (s): We compute the prefix sum array as follows:
- Initialize by setting s[0] = 0
- For i=1, add s[0] + nums[0] -> s[1] = 0 + 1 = 1
- For i=2, add s[1] + nums[1] -> s[2] = 1 + 3 = 4
- For i=3, add s[2] + nums[2] -> s[3] = 4 + 4 = 8
- For i=4, add s[3] + nums[3] -> s[4] = 8 + 9 = 17
The prefix sum array now is s = [0, 1, 4, 8, 17].
Binary Search: We apply binary search to find the split point for query 6: Using bisect_left, we find that:
- The index i for the first element greater than 6 is 3 (nums[3] = 9).
- The index j for the first element equal to or greater than 6 is also 3.
Calculating the Number of Operations:
- For elements greater than 6 (nums[3]), calculate the reduction operations: The sum of elements bigger than 6 is s[-1] - s[i] = 17 - 8 = 9. Total reduction needed (len(nums) - i) * 6 = (4 - 3) * 6 = 6. Total decrease operations t = 6 - 9 = -3 (since we are decreasing, this number should be positive, so we take absolute: 3).
- For elements less than 6 (from nums[0] to nums[2]), calculate the increase operations: Total increase needed is 6 * i = 6 * 3 = 18. Sum of these elements is s[i] = 8. Total increase operations needed t = 18 - 8 = 10.
Consequently, the total minimum operations for 6 is the sum of increase and decrease operations: 3 + 10 = 13.
Build the Result Array (ans): Append 13 to the result array ans. Since there is only one query value, ans = [13].

By following these steps, the final answer we get is that 13 operations are necessary to make all elements in the array nums equal to the single query value 6.

Solution Implementation

1from bisect import bisect_left
2from itertools import accumulate
3
4class Solution:
5    def minOperations(self, nums, queries):
6        # Sort the input array to allow binary search operations 
7        nums.sort()
8      
9        # Accumulate the sum of elements in the array, starting from 0
10        prefix_sum = list(accumulate(nums, initial=0))
11      
12        # Initialize the list to store results of each query
13        answer = []
14      
15        # Evaluate each query in the queries list
16        for x in queries:
17            # Find the index of the smallest number in nums that is greater than x
18            index = bisect_left(nums, x + 1)
19            # Calculate the total operations needed for the larger portion of nums
20            total_ops = prefix_sum[-1] - prefix_sum[index] - ((len(nums) - index) * x)
21          
22            # Find the index of the smallest number in nums that is greater than or equal to x
23            index = bisect_left(nums, x)
24            # Add the operations needed for the smaller portion of nums
25            total_ops += x * index - prefix_sum[index]
26          
27            # Append the total number of operations to answer
28            answer.append(total_ops)
29      
30        # Return the final result after processing all queries
31        return answer
32

1import java.util.ArrayList;
2import java.util.Arrays;
3import java.util.List;
4
5class Solution {
6    public List<Long> minOperations(int[] nums, int[] queries) {
7        // Sort the array of numbers in non-decreasing order
8        Arrays.sort(nums);
9      
10        int n = nums.length; // Number of elements in nums
11        // Create a prefix sum array with an extra space for ease of calculations
12        long[] prefixSum = new long[n + 1];
13      
14        // Calculate the prefix sum
15        for (int i = 0; i < n; ++i) {
16            prefixSum[i + 1] = prefixSum[i] + nums[i];
17        }
18      
19        // Initialize the answer list
20        List<Long> ans = new ArrayList<>();
21      
22        // Process each query
23        for (int x : queries) {
24            // Search for the first element in nums that is greater than or equal to x + 1
25            int index = binarySearch(nums, x + 1);
26            // Calculate the total sum needed to reduce elements larger or equal to x + 1 to x
27            long totalOperations = prefixSum[n] - prefixSum[index] - 1L * (n - index) * x;
28          
29            // Search for the first element in nums that is greater than or equal to x
30            index = binarySearch(nums, x);
31            // Add the total sum needed to increase elements smaller than x to x
32            totalOperations += 1L * x * index - prefixSum[index];
33          
34            // Add the calculated operations for the current query to the answer list
35            ans.add(totalOperations);
36        }
37      
38        return ans; // Return the answer list
39    }
40
41    // Helper method to perform binary search
42    private int binarySearch(int[] nums, int x) {
43        int left = 0;
44        int right = nums.length;
45        // Binary search to find the index of the first number greater or equal to x
46        while (left < right) {
47            int mid = (left + right) >> 1; // Middle position
48            if (nums[mid] >= x) {
49                right = mid;
50            } else {
51                left = mid + 1;
52            }
53        }
54        return left; // Return the first index where nums[index] >= x
55    }
56}
57

1#include <vector>
2#include <algorithm>
3
4class Solution {
5public:
6    // Function to find the minimum operations needed for each query
7    vector<long long> minOperations(vector<int>& nums, vector<int>& queries) {
8        // Sort the input numbers array for binary search
9        sort(nums.begin(), nums.end());
10      
11        // Calculate the size of the nums array
12        int numsSize = nums.size();
13      
14        // Prefix sum array initialized with an extra element for ease of calculations
15        vector<long long> prefixSums(numsSize + 1, 0);
16      
17        // Compute the prefix sums for efficient range sum queries
18        for (int i = 0; i < numsSize; ++i) {
19            prefixSums[i + 1] = prefixSums[i] + nums[i];
20        }
21      
22        // Vector to store the answer for each query
23        vector<long long> answer;
24      
25        // Iterate over all queries to compute the minimum operations
26        for (auto& queryValue : queries) {
27            // Find the first element greater than the query value
28            int upperIndex = lower_bound(nums.begin(), nums.end(), queryValue + 1) - nums.begin();
29          
30            // Calculate the operations needed for transforming numbers greater than the query value
31            long long operationsForGreater = prefixSums[numsSize] - prefixSums[upperIndex] - 1LL * (numsSize - upperIndex) * queryValue;
32          
33            // Find the first element that is not less than the query value
34            int lowerIndex = lower_bound(nums.begin(), nums.end(), queryValue) - nums.begin();
35          
36            // Calculate the operations needed for transforming numbers less than or equal to the query value
37            long long operationsForLessOrEqual = 1LL * queryValue * lowerIndex - prefixSums[lowerIndex];
38          
39            // The total operations is the sum of operations for both ranges
40            long long totalOperations = operationsForGreater + operationsForLessOrEqual;
41          
42            // Append the total operations to the answer vector
43            answer.push_back(totalOperations);
44        }
45      
46        // Return the final answer vector
47        return answer;
48    }
49};
50

1// Function to calculate the minimum operations required for the queries
2function minOperations(nums: number[], queries: number[]): number[] {
3    // Sort the array in non-decreasing order
4    nums.sort((a, b) => a - b);
5
6    const length = nums.length;
7
8    // Precompute the prefix sums of the sorted array
9    const prefixSums: number[] = new Array(length + 1).fill(0);
10    for (let i = 0; i < length; ++i) {
11        prefixSums[i + 1] = prefixSums[i] + nums[i];
12    }
13
14    // Binary search function to find the first index where value is >= x
15    const binarySearch = (x: number): number => {
16        let left = 0;
17        let right = length;
18        while (left < right) {
19            const mid = (left + right) >> 1; // Equivalent to Math.floor((left + right) / 2)
20            if (nums[mid] >= x) {
21                right = mid;
22            } else {
23                left = mid + 1;
24            }
25        }
26        return left;
27    };
28
29    const answers: number[] = [];
30
31    // Process each query
32    for (const query of queries) {
33        // Find the first number greater than query
34        const greaterIndex = binarySearch(query + 1);
35
36        // Calculate the total of subtracting query from elements greater than query
37        let total = prefixSums[length] - prefixSums[greaterIndex] - (length - greaterIndex) * query;
38
39        // Find the first number greater than or equal to query
40        const equalOrGreaterIndex = binarySearch(query);
41
42        // Calculate the total of adding query to elements less than query
43        total += query * equalOrGreaterIndex - prefixSums[equalOrGreaterIndex];
44
45        // Add the result to the answer array
46        answers.push(total);
47    }
48
49    return answers;
50}
51

Time and Space Complexity

The given algorithm involves sorting the array nums and performing binary searches for each query in queries.

Time Complexity

Sorting the nums array has a time complexity of O(n \log n) where n is the length of the nums array.
The accumulate function is linear, contributing O(n) to the time complexity.
For each query, a binary search is performed twice using bisect_left, which has a time complexity of O(\log n). Since there are q queries, the total complexity for this part is O(q \log n).

The total time complexity combines these contributions, resulting in O(n \log n) + O(n) + O(q \log n). Since O(n \log n) dominates O(n), and the number of queries q could vary independently of n, the overall time complexity is O((n + q) \log n).

However, since the reference answer only specifies O(n \log n), it implies that n is the dominant term, and q is expected to be not significantly larger than n.

Space Complexity

The sorted array is a modification in place, so it does not add to the space complexity.
The s variable is a list storing the cumulative sum, contributing O(n) to the space complexity.
A constant amount of extra space is used for variables i, t, and the iterative variables, which does not depend on the size of the input.

Thus, the overall space complexity is O(n).

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

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