从列表中原位删除部分元素的正确方法

532 阅读4分钟
原文链接: zhuanlan.zhihu.com

(图文无关)

这是一篇面向编程初学者的文章。看到 挑战自己|LeetCode 刷题开胃菜 - 知乎专栏,我觉得有必要讲一下这几道LeetCode的正确做法。

简单来说,原题中共同的要求是给定一个数组(Python列表),在不构建新的数组的情况下,移除数组中某些符合条件的元素。比如说,数组[1,2,3,1,3,4],移除所有的1之后,就变成了[2,3,3,4]。可以注意到,列表在移除某些元素之后,长度会变短,元素的下标会做相应的调整。如果是世界最好的语言PHP,它的数组其实是数字做key的hashmap,那么只要删掉相应的key,然后用array_merge重整一下就行了(当然,这种做法完全不符合题意的in place);但对Python来说,它的list是真正的顺序表,需要仔细考虑删除的过程。我知道大家都喜欢[x for x in array if ...],但是题意是要用原始列表,那么就应该按题目要求。

为了简单,我们首先将问题简化一下:我们有一个列表,其中某些元素是None,现在我们要操作这个列表使其中仅有非None的元素留下来。显然,其他问题可以通过先将要删除的元素设置为None(就像javascript里面的delete a[index]那样,在数组里留个洞),然后调用这个过程来实现。

当然,最直接的方法,我们可以选择循环调用remove(None),直到抛出异常。不过这种做法相对于其他语言来说太过于作弊了,而且性能上也不是最好的(后面会分析),暂时不提。

一种常见的错误做法是这样的:

a = [1, 1, None, None, 3, 3, 4, 4]
for i in range(0, len(a)):
    if i >= len(a):
        break
    if a[i] is None:
        del a[i]

运行会发现得到的结果是不正确的,有一个None没有删掉。为什么呢?程序中遇到i = 2的时候,调用del a[2],会将后面的元素前移一位,于是之间下标为3的元素会变成下标为2的元素,而下一次处理的是i = 3,于是就漏掉了一个元素。

由于删除一个元素的时候,这个元素后面的元素下标都会改变,一种简单的想法就是反过来从后向前删除:

a = [1, 1, None, None, 3, 3, 4, 4]
for i in range(len(a) - 1, -1, -1):
    if a[i] is None:
        del a[i]

这种做法是正确的。但是,如果数组很大,这种做法会比较慢。原因在于,删除一个元素的时候,会重新复制这之后所有的元素,这个过程需要的时间跟数组长度成正比,这样最坏情况下,这个程序就需要O(n^2)的时间来完成这个过程(想象一个1和None间隔的数组)。O(f(n))这个符号表示渐进复杂度,当问题规模用n来表示(对这个问题来说,就是数组长度为n)时,需要消耗的时间的增长速度的上界不超过f(n)的某个较大的倍数。

我们希望对于长度为n的数组,只使用O(n)的操作来完成这个过程。实际上,相应的代码是很简单的:

j = 0
for i in range(0, len(a)):
    if a[i] is not None:
        a[j] = a[i]
        j += 1
del a[j:]

解释下原理,我们考虑去掉None元素的这个过程,它的原理在于将非None的元素从原来的位置上,调整到更靠前的位置上。在这个过程中,元素的下标一定不会增加,而是会不变或者减少。因此我们读取的下一个元素一定还没有被修改过,那么我们用i来表示准备移动的元素的下标,用j来表示下一个要放入的下标位置,就写出了上面的代码。

再来考虑最初的问题,我们要移除数组中某些符合条件的数,当然没有必要真的将它设置为None,只需要在读取的过程中,一边读取一边判断就好了:

def filter_in_place(array, filter_):
    j = 0
    for i in range(0, len(array)):
        if filter_(array[i]):
            array[j] = array[i]
            j += 1
    del array[j:]
    return len(array)

删除某个值的元素就可以写成:

def remove_in_place(array, num):
    filter_in_place(array, lambda x: x == num)

删除重复元素稍微复杂些,要点在于列表已经排好了序这件事。可以构造一个闭包来完成这件事:

last_num = None
def filter_duplicate(x):
    nonlocal last_num
    if x == last_num:
        return False
    else:
        last_num = x
        return True

def remove_duplicate_in_place(array):
    return filter_in_place(array, filter_duplicate)

如果Python2的话可以用一个可变对象来代替nonlocal:

last_num = [None]
def filter_duplicate(x):
    if x == last_num[0]:
        return False
    else:
        last_num[0] = x
        return True

def remove_duplicate_in_place(array):
    return filter_in_place(array, filter_duplicate)

用一个类也是一样的:

class DuplicatedFilter(object):
    def __init__(self):
        self.last_num = None

    def __call__(self, x):
        if x == self.last_num:
            return False
        else:
            self.last_num = x
            return True

def remove_duplicate_in_place(array):
    return filter_in_place(array, DuplicatedFilter())

那最后,要保留两个,也很简单了:

class DuplicatedFilter(object):
    def __init__(self):
        self.last_num = None
        self.count = 0

    def __call__(self, x):
        if x == self.last_num:
            if self.count >= 1:
                return False
            else:
                self.count += 1
                return True
        else:
            self.last_num = x
            self.count = 0
            return True

def remove_duplicate_in_place(array):
    return filter_in_place(array, DuplicatedFilter())

请自行写出闭包的版本作为课后练习。