Python中 shift 和 merge 列表元素的最高效方法 (2048)

Question

我有一个函数，主要是把一个列表的内容右对齐，同时把两个相同的元素合并成一个（这个列表至少会有一个元素）：

def shift(sequence):
    for i in range(len(sequence)-1):
        current_value = sequence[i]
        next_value = sequence[i+1]
        if next_value == 0:
            sequence[i], sequence[i+1] = 0, current_value
        elif current_value == next_value:
            sequence[i], sequence[i+1] = 0, current_value*2
    return sequence

这里有一些示例输入和输出：

>>> shift([0, 0, 1, 0])
[0, 0, 0, 1]
>>> shift([1, 1, 0, 0])
[0, 0, 0, 2]
>>> shift([2, 0, 1, 0])
[0, 0, 2, 1]

那么，最有效的方法是什么呢？如果这个操作是对一个矩阵的每一行进行的，有没有比这样做更有效的方法：

matrix = [shift(row) for row in matrix]

另外，如果我想把矩阵向其他三个方向移动（除了向右），有没有比这三种方法更有效的方式：

#Left
matrix = [shift(row[::-1])[::-1] for row in matrix]

#Down
matrix = map(list, zip(*[shift(row) for row in map(list, zip(*matrix))]))

#Up
matrix = map(list, zip(*[shift(row[::-1])[::-1] for row in map(list, zip(*matrix))]))

如果这些移动操作是重复进行的（而且每次矩阵中的某个值都会改变），有没有什么需要注意的地方，以提高效率？

编辑

我的函数并不是总能正常工作：

>>> shift([1, 1, 1, 1])
[0, 2, 0, 2]

输出应该是：

[0, 0, 2, 2]

更多的预期输入和输出：

[1, 1, 1]             --> [0, 1, 2]
[1, 2, 2, 3, 5, 5, 2] --> [0, 0, 1, 4, 3, 10, 2]

编辑 2

移动不一定要向右，也可以向其他方向。

Answer 1

这是我目前的解决方案。它比Kirk Strauser的方案快大约60%。

def shift(self, sequence, right=False):
    if right:
        sequence = sequence[::-1]
    values = []
    empty = 0
    for n in sequence:
        if values and n == values[-1]:
            values[-1] = 2*n
            empty += 1
        elif n:
            values.append(n)
        else:
            empty += 1
    values += [0]*empty
    if right:
        values = values[::-1]
    return values

这是我更高效的方案：

def shift2(length, sequence, right=False):
    if right:
        sequence = sequence[::-1]
    values = [0]*length
    full = 0
    for n in sequence:
        if full and n == values[full]:
            values[full] = 2*n
        elif n:
            values[full] = n
            full += 1
    if right:
        values = values[::-1]
    return values

Kirk的方案：

def streaming_sum(sequence):
    values = reversed(sequence)
    last = next(values)
    for value in values:
        if value == last:
            yield last + value
            last = 0
        else:
            yield last
            last = value
    yield last

def shift2(sequence):
    length = len(sequence)
    reduced = list(reversed(list(filter(None, streaming_sum(sequence)))))
    return [0] * (length - len(reduced)) + reduced

这是我对Kirk的函数进行改进后的版本（速度提升40%）：

def shift3(sequence):
    length = len(sequence)
    reduced = [n for n in filter(None, streaming_sum(sequence))][::-1]
    return [0] * (length - len(reduced)) + reduced

---

计时：

from timeit import Timer

tests = [[1000, 1000, 1000, 1000],
         [1000, 0, 0, 1000],
         [0, 1000],
         [1000, 1000, 0, 500, 0, 500, 1000, 0, 0, 100, 100, 100]]
t1, t2, t3 = 0, 0, 0

for test in tests:
    t1 += Timer(lambda: shift(test)).timeit()
    t2 += Timer(lambda: shift2(test)).timeit()
    t3 += Timer(lambda: shift3(test)).timeit()

>>> print(t1, t2, t3)
10.706327316242147 26.92895738572211 16.65189852514444

我这个方案的计时结果与我更高效的方案相比，以及没有右对齐选项的情况下的计时结果：

10.502816527107633 8.503653343656246 8.15101370397509

Answer 2

这是否更有效率就要看 timeit 的结果了：

def streaming_sum(sequence):
    values = reversed(sequence)
    last = next(values)
    for value in values:
        if value == last:
            yield last + value
            last = 0
        else:
            yield last
            last = value
    yield last


def shift(sequence):
    length = len(sequence)
    reduced = list(reversed(filter(None, streaming_sum(sequence))))
    return [0] * (length - len(reduced)) + reduced


for sequence, expected in [
    ([0, 0, 1, 0], [0, 0, 0, 1]),
    ([1, 1, 0, 0], [0, 0, 0, 2]),
    ([2, 0, 1, 0], [0, 0, 2, 1]),
    ([1, 1, 1, 1], [0, 0, 2, 2]),
    ([1, 1, 1], [0, 1, 2]),
    ([1, 2, 2, 3, 5, 5, 2], [0, 0, 1, 4, 3, 10, 2]),
]:
    actual = shift(sequence)
    assert actual == expected, (actual, expected)