leetcode算法题：100277. 使数组中位数等于 K 的最少操作数

写在前面

今天上午做leetcode周赛，第三题是：100277. 使数组中位数等于 K 的最少操作数

其实当时已经做的差不多了，但是有2个测试用例未通过，报超时了。事后看题解，思路根别人是一样的，但是别人用到的排序算法是go自带的sort.Ints，而我用的是自己写的一个快速排序算法。我将排序算法改成sort.Ints就通过了。

感觉有点神奇，不知道go自带的sort.Ints算法是如何实现的，比排序算法还快。本文记录一下。

思路

代码

下面直接上作者做题的代码

func MinOperationsToMakeMedianK(nums []int, k int) int64 {
	n := len(nums)
	mid := n / 2
	fastSort(nums, 0, n-1)
	//sort.Ints(nums)
	result := int64(0)
	if nums[mid] < k {
		// 需要将中位数右边的数字调整
		for i := mid; i < n; i++ {
			if nums[i] < k {
				result += int64(k - nums[i])
			} else {
				break
			}
		}
	} else {
		// 需要将中位数左边的数字调整
		for i := mid; i >= 0; i-- {
			if nums[i] > k {
				result += int64(nums[i] - k)
			} else {
				break
			}
		}
	}

	return result
}

// 快速排序
func fastSort(data []int, begin, end int) {
	if begin >= end {
		return
	}
	i, j := begin, end
	for i < j {
		for i < j && data[i] <= data[j] {
			i++
		}
		if i >= j {
			break
		}
		data[i], data[j] = data[j], data[i]

		for i < j && data[i] <= data[j] {
			j--
		}
		if i >= j {
			break
		}
		data[i], data[j] = data[j], data[i]
	}
	fastSort(data, i+1, end)
	fastSort(data, begin, i-1)
}

分析

上面这段代码，有3个测试用例不通过

上面这个测试用例里面很多1，相当于是已经排序好的，无需排序。而快速排序对已经排序好的数组（递增或者递减）时间复杂的是最大的，达到O(n^2)。

而go自带的sort.Ints算法不知道是如何实现的，用它就能通过这个测试用例。

优化

针对上述测试用例，快速排序能否优化一下呢？上述测试用例的特征是很多相同的元素。

于是作者进行了如下优化

// 快速排序
func fastSort(data []int, begin, end int) {
	if begin >= end {
		return
	}
	i, j := begin, end
	for i < j {
		for i < j && data[i] <= data[j] {
			i++
		}
		if i >= j {
			break
		}
		data[i], data[j] = data[j], data[i]

		for i < j && data[i] <= data[j] {
			j--
		}
		if i >= j {
			break
		}
		data[i], data[j] = data[j], data[i]
	}
	newBegin := i
	for newBegin < end && data[newBegin] == data[newBegin+1] {
		newBegin++
	}
	newEnd := i
	for newEnd > begin && data[newEnd] == data[newEnd-1] {
		newEnd--
	}
	fastSort(data, newBegin+1, end)
	fastSort(data, begin, newEnd-1)
}

在递归调用fastSort之前，判断一下，对于递归右边子数组，开始位置如果跟i位置相同则跳过，左边同样处理。这样处理之后本题就通过了。