2602. Minimum Operations to Make All Array Elements Equal

Medium

You are given an array nums consisting of positive integers.

You are also given an integer array queries of size m. For the ith query, you want to make all of the elements of nums equal to queries[i]. You can perform the following operation on the array any number of times:

• Increase or decrease an element of the array by 1.

Return an array answer of size m where answer[i] is the minimum number of operations to make all elements of nums equal to queries[i].

Note that after each query the array is reset to its original state.

Example 1:

Input: nums = [3,1,6,8], queries = [1,5]
Output: [14,10]
Explanation: For the first query we can do the following operations:
- Decrease nums[0] 2 times, so that nums = [1,1,6,8].
- Decrease nums[2] 5 times, so that nums = [1,1,1,8].
- Decrease nums[3] 7 times, so that nums = [1,1,1,1].
So the total number of operations for the first query is 2 + 5 + 7 = 14.
For the second query we can do the following operations:
- Increase nums[0] 2 times, so that nums = [5,1,6,8].
- Increase nums[1] 4 times, so that nums = [5,5,6,8].
- Decrease nums[2] 1 time, so that nums = [5,5,5,8].
- Decrease nums[3] 3 times, so that nums = [5,5,5,5].
So the total number of operations for the second query is 2 + 4 + 1 + 3 = 10.

Example 2:

Input: nums = [2,9,6,3], queries = [10]
Output: [20]
Explanation: We can increase each value in the array to 10. The total number of operations will be 8 + 1 + 4 + 7 = 20.

Constraints:

• n == nums.length

• m == queries.length

• 1 <= n, m <= 10^5

• 1 <= nums[i], queries[i] <= 10^9

解題

func minOperations(nums []int, queries []int) []int64 {
	l := len(nums)
	sort.Ints(nums) // 排序 nums，排序完的 nums 是 [1 3 6 8]

	prefixSums := make([]int,l) // 前綴和，以範例一為例，會是 [1 4 10 18]
	prefixSums[0] = nums[0]
	for i:=1; i<l; i++ {
		prefixSums[i] = prefixSums[i-1]+nums[i]
	}

	result := make([]int64, len(queries))
	for i, query := range queries {
		idx := sort.SearchInts(nums, query)

		if idx == 0 {
			result[i] = int64(prefixSums[l-1] - query*l)
		} else {
			// 比 query 大的數字，需要做 Decrease，也就是最後的前綴和 - idx 前的前綴和，最後扣掉的 query*(l-idx) 代表要留下的部分，idx 後的每個數字都會變成 query
			result[i] = int64(prefixSums[l-1] - prefixSums[idx-1] - query*(l-idx)) 
			// 比 query 小的數字，需要做 Increase，需要增加的次數就是 query*idx - prefixSums[idx-1] 
			result[i] += int64(query*idx - prefixSums[idx-1])
		}
	}

	return result
}