方法 2:借助快速排序 partition 过程的一趟扫描法
我们先来回顾快速排序 partition
过程:随机选择一个元素作为切分元素(pivot
),然后经过一次扫描,通过交换元素的位置使得数组按照数值大小分成以下 3 个部分:
1、严格小于 pivot
;
2、等于 pivot
;
3、严格大于 pivot
。
根据这个操作的启发,我们也完全可以设计一个算法使用扫描一次的办法把只含有 0
、1
、2
的整数数组整理成 0
、1
、2
紧凑排列的样子。
而写对这个算法,我们需要借助“循环不变式(loop invariant)”这个概念。循环不变式是程序员根据问题的特点构造的一个断言,也就是对即将在循环中完成的事情的一个描述,这个断言在循环的开始、循环的运行过程中、循环的结束时都保持为真。循环不变式用于证明算法的正确性。
由于我们需要把数组分成三个部分,为此我需要设置两个变量作为不同部分的分界,和一个循环变量 i
,于是循环不变量可以这样定义的:
- 所有在子区间
[0, zero)
的元素都等于0
; - 所有在子区间
[zero, i)
的元素都等于1
; - 所有在子区间
[two, len - 1]
的元素都等于2
。
(这里要画图。)
说明:设计循环不变量的原则是“不重不漏”。
1、len
是数组的长度;
2、变量 zero
是前两个子区间的分界点,一个是闭区间,另一个就必须是开区间;
3、变量 i
是循环变量,一般设置为开区间,表示 i
之前的元素是遍历过的;
4、two
是另一个分界线,我设计成闭区间。
于是编码要解决以下三个问题:
1、变量初始化应该如何定义;
2、在遍历的时候,zero
和 two
应该怎样移动,是先加减还是先交换;
3、什么时候终止循环。
而解决这三个问题,完全看循环不变量的定义。
- 编码的时候,
zero
和two
初始化的值就应该保证上面的三个子区间全为空; - 在遍历的过程中,“索引先加减再交换”、还是“先交换再加减”就看初始化的时候变量在哪里;
- 退出循环的条件也看上面定义的循环不变量,在
i == two
成立的时候,上面的三个子区间就正好“不重不漏”地覆盖了整个数组,并且给出的性质成立,题目的任务也就完成了。
以下给出不同的写法,循环不变量的定义写在了注释中。使用“循环不变量”编码是为了便于我们处理细节,也方便别人看懂你的算法。
参考代码 1:
Java 代码:
import java.util.Arrays;
public class Solution {
public void sortColors(int[] nums) {
int len = nums.length;
if (len < 2) {
return;
}
// all in [0, zero) = 0
// all in [zero, i) = 1
// all in [two, len - 1] = 2
// 循环终止条件是 i == two,那么循环可以继续的条件是 i < two
// 为了保证初始化的时候 [0, zero) 为空,设置 zero = 0,
// 所以下面遍历到 0 的时候,先交换,再加
int zero = 0;
// 为了保证初始化的时候 [two, len - 1] 为空,设置 two = len
// 所以下面遍历到 2 的时候,先减,再交换
int two = len;
int i = 0;
// 当 i == two 上面的三个子区间正好覆盖了全部数组
// 因此,循环可以继续的条件是 i < two
while (i < two) {
if (nums[i] == 0) {
swap(nums, i, zero);
zero++;
i++;
} else if (nums[i] == 1) {
i++;
} else {
two--;
swap(nums, i, two);
}
}
}
private void swap(int[] nums, int index1, int index2) {
int temp = nums[index1];
nums[index1] = nums[index2];
nums[index2] = temp;
}
}
参考代码 2:
Java 代码:
public class Solution {
public void sortColors(int[] nums) {
int len = nums.length;
if (len < 2) {
return;
}
// all in [0, zero] = 0
// all in (zero, i) = 1
// all in (two, len - 1] = 2
// 为了保证初始化的时候 [0, zero] 为空,设置 zero = -1,
// 所以下面遍历到 0 的时候,先加,再交换
int zero = -1;
// 为了保证初始化的时候 (two, len - 1] 为空,设置 two = len - 1
// 所以下面遍历到 2 的时候,先交换,再减
int two = len - 1;
int i = 0;
// 当 i == two 的时候,还有一个元素还没有看,
// 因此,循环可以继续的条件是 i <= two
while (i <= two) {
if (nums[i] == 0) {
zero++;
swap(nums, i, zero);
i++;
} else if (nums[i] == 1) {
i++;
} else {
swap(nums, i, two);
two--;
}
}
}
private void swap(int[] nums, int index1, int index2) {
int temp = nums[index1];
nums[index1] = nums[index2];
nums[index2] = temp;
}
}