LeetCode 2161：按指定基准值划分数组的优化解法问询

问题回顾 #

给定0索引的整数数组nums和整数pivot，重排数组需满足：

• 所有小于pivot的元素在大于pivot的元素之前
• 等于pivot的元素位于两者中间
• 小于和大于pivot的元素相对顺序保持不变

示例：输入nums = [9,12,5,10,14,3,10]，pivot = 10，输出[9,5,3,10,10,12,14]

你当前实现了类似插入排序的O(n²)时间、O(1)空间解法，以下是针对你问题的解答：

1. 时间复杂度优于O(n²)的解法（O(n)时间，O(n)空间） #

这是最直观的优化思路，通过分三个区域存储元素再拼接，既能保证相对顺序，又能将时间复杂度降到O(n)：

public int[] pivotArray(int[] nums, int pivot) {
    int n = nums.length;
    int[] result = new int[n];
    int idx = 0;
    
    // 先存入所有小于pivot的元素
    for (int num : nums) {
        if (num < pivot) {
            result[idx++] = num;
        }
    }
    // 再存入所有等于pivot的元素
    for (int num : nums) {
        if (num == pivot) {
            result[idx++] = num;
        }
    }
    // 最后存入所有大于pivot的元素
    for (int num : nums) {
        if (num > pivot) {
            result[idx++] = num;
        }
    }
    
    return result;
}

该解法三次遍历数组，时间复杂度O(n)，空间复杂度O(n)，代码简洁且严格满足相对顺序要求。

2. O(n)时间且O(1)空间的原地解法 #

要实现原地O(n)时间重排，需注意必须保留小于、大于pivot元素的相对顺序（这和快速排序的分区逻辑不同，快排不保证相对顺序），可以通过两次遍历完成：

实现步骤 #

1. 第一步：将所有小于pivot的元素移到数组前部

   • 用指针left标记下一个小于pivot元素的存放位置
   • 遍历数组，遇到小于pivot的元素就与left位置元素交换，随后left右移
   • 完成后，left之前的元素均为小于pivot的元素，且相对顺序不变

2. 第二步：将所有等于pivot的元素移到小于区域之后

   • 用指针mid从left位置开始，标记下一个等于pivot元素的存放位置
   • 从left开始遍历数组，遇到等于pivot的元素就与mid位置元素交换，随后mid右移
   • 完成后，left到mid-1之间的元素均为等于pivot的元素，剩余元素则为大于pivot的元素，且相对顺序不变

代码实现 #

public int[] pivotArray(int[] nums, int pivot) {
    int n = nums.length;
    int left = 0;
    
    // 移动小于pivot的元素到前部
    for (int i = 0; i < n; i++) {
        if (nums[i] < pivot) {
            int temp = nums[i];
            nums[i] = nums[left];
            nums[left] = temp;
            left++;
        }
    }
    
    // 移动等于pivot的元素到小于区域之后
    int mid = left;
    for (int i = left; i < n; i++) {
        if (nums[i] == pivot) {
            int temp = nums[i];
            nums[i] = nums[mid];
            nums[mid] = temp;
            mid++;
        }
    }
    
    return nums;
}

该解法为原地操作，空间复杂度O(1)，时间复杂度O(n)，完全符合题目所有要求。

对原解法的补充说明 #

你原有的插入排序式思路，虽然实现了O(1)空间，但最坏情况下（例如数组全为大于pivot的元素，仅最后一个元素小于pivot），每次交换都需要遍历前面所有元素，时间复杂度会达到O(n²)，对于大规模数组的处理效率较低。

内容的提问来源于stack exchange，提问作者Some Name