数组三分

1 相关问题描述

1.1 荷兰国旗问题^[1]

荷兰国旗问题是一个经典的计算机科学问题，其来源于荷兰国旗的三色（红、白、蓝）设计，要求通过一次遍历将包含三种颜色的数组按颜色进行排序。数组中的每个元素代表一种颜色，分别用 0、1 和 2 表示红色、白色和蓝色。目标是将这个数组重新排列，使得所有的 0 都位于数组的前面，所有的 1 都位于中间，所有的 2 都位于数组的后面。

1.2 三路快速排序

三路快速排序（3-way QuickSort）是快速排序算法的一种变体，特别适合处理有大量重复元素的数组。在标准快速排序中，我们选择一个元素作为基准（pivot），然后将数组分为两部分：小于基准的元素和大于基准的元素。而在三路快速排序中，数组被分为三个部分：小于基准的元素、等于基准的元素和大于基准的元素。这样，每次递归时，等于基准的元素已经被正确排序，因此下一次递归只需要处理小于和大于基准的两个子数组。

三路快速排序的主要优点是它能够减少对重复元素的无谓比较，从而提高排序的效率。特别是当数组中存在大量重复元素时，与传统的快速排序相比，三路快速排序可以显著减少递归的深度，提高排序速度。

1.3 三路快速选择

快速选择算法^[2]是一种高效地找到数组中第 k 小或第 k 大元素的方法，它是快速排序的一个变种。它通过选择一个基准元素（pivot），把数组分为小于和大于基准的两部分来工作。与快速排序不同，快速选择只对含有第 k 小元素的子数组进行递归，减少了计算量。

三路快速选择是快速选择的扩展，适用于含有大量重复元素的数组。它按小于、等于、大于基准元素的三部分来分数组，并能在寻找第 k 小或大元素时跳过重复项。这种方法有效减少了重复元素的比较次数，提升了效率。

三路快速选择在处理重复元素多的数组时更加高效，减少了不必要的比较和递归深度，提高了性能。

例如 LeetCode 215. 数组中的第K个最大元素^[3] 中，有一个测试用例为:

nums = [1, 2, 3, 4, 5] + [1] * 99000 + [-5, -4, -3, -2, -1]
k = 50000

如果使用普通的快速选择算法，极易发生栈溢出。

2 代码实现

def partition(nums: List[int], pivot:int):
    """
    在一次遍历中将列表 nums 中小于、等于、大于 pivot 的元素
    分别移动到列表的左侧、中间、右侧，使用三指针的方法。
    """
    # i 用于遍历 nums 中的每个元素
    # left 指向当前未处理的最左侧元素的位置，用于放置小于 pivot 的元素
    # right 指向当前未处理的最右侧元素的位置，用于放置大于 pivot 的元素
    i, left, right = 0, 0, len(nums) - 1

    while i <= right:
        if nums[i] < pivot:
            # 如果当前元素小于 pivot，则将其与 left 指针所指元素交换
            # 交换后，left 和 i 都向右移动一位
            nums[i], nums[left] = nums[left], nums[i]
            left += 1
            i += 1
        elif nums[i] > pivot:
            # 如果当前元素大于 pivot，则将其与 right 指针所指元素交换
            # 交换后，right 向左移动一位，而 i 不变，因为交换后的新元素还未被比较
            nums[i], nums[right] = nums[right], nums[i]
            right -= 1
        else:
            # 如果当前元素等于 pivot，则不需要交换
            # i 向右移动一位，继续遍历
            i += 1

首先理解一下各个区域的含义

当 nums[i] < pivot 时：

if nums[i] < pivot:
    # 如果当前元素小于 pivot，则将其与 left 指针所指元素交换
    # 交换后，left 和 i 都向右移动一位
    nums[i], nums[left] = nums[left], nums[i]
    left += 1
    i += 1

当 nums[i] = pivot 时：

else:
    # 如果当前元素等于 pivot，则不需要交换
    # i 向右移动一位，继续遍历
    i += 1

当 nums[i] > pivot 时：

elif nums[i] > pivot:
    # 如果当前元素大于 pivot，则将其与 right 指针所指元素交换
    # 交换后，right 向左移动一位，而 i 不变，因为交换后的新元素还未被比较
    nums[i], nums[right] = nums[right], nums[i]
    right -= 1

3 引用

1. LeetCode 75. 颜色分类 ↩
2. 维基百科：快速选择算法 ↩
3. LeetCode 215. 数组中的第K个最大元素 ↩