python - 在 Python 中生成和压缩两个列表的最干净有效的方法

Question

给定这两个列表

zeros = [2,3,1,2]
ones = [3,4,5]

(条件总是 len(zeros) == len(ones) + 1)

我想创建一个列表，其中交替包含列表中提到的大小的 0 和 1。我可以通过以下方式实现:

zeros_list = [[0]*n for n in zeros]
ones_list = [[1]*n for n in ones]
output = [z for x in zip(zeros_list, ones_list) for y in x for z in y]
output += [0]*zeros[-1]
print(output)
> [0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 0]

然而，这是最有效/最干净的方式吗？我得到了 2.66 µs ± 78.8 ns 的性能，但我仍然认为这可以在单行中完成，而且效率可能更高

Answer 1

两个明显更快的解决方案，使用“技巧”对第一个 零而不是最后一个零进行特殊处理，使用它来初始化输出。

def superb_rain(zeros, ones):
    zeros = iter(zeros)
    output = [0] * next(zeros)
    for o in ones:
        output += (1,) * o
        output += (0,) * next(zeros)
    return output

def superb_rain2(zeros, ones):
    z = iter(zeros).__next__
    output = [0] * z()
    for o in ones:
        output += (1,) * o
        output += (0,) * z()
    return output

(正如@schwobaseggl 指出的那样，元组使速度提高了大约 30%。)

基准测试结果:

0.14 us  0.13 us  0.13 us  baseline
3.04 us  3.02 us  2.98 us  original
3.27 us  3.19 us  3.29 us  chepner_1
5.03 us  5.12 us  5.25 us  chepner_2
4.66 us  4.74 us  4.68 us  chepner_2__superb_rain
2.52 us  2.53 us  2.47 us  Alain_T
3.35 us  3.27 us  3.42 us  python_user
1.02 us  0.99 us  1.04 us  superb_rain
1.07 us  1.11 us  1.09 us  superb_rain2

基准代码:

import timeit
from itertools import zip_longest, cycle, islice, repeat, chain

def baseline(zeros, ones):
    pass

def original(zeros, ones):
    zeros_list = [[0]*n for n in zeros]
    ones_list = [[1]*n for n in ones]
    output = [z for x in zip(zeros_list, ones_list) for y in x for z in y]
    output += [0]*zeros[-1]
    return output

def chepner_1(zeros, ones):
    return list(chain.from_iterable(chain.from_iterable(zip_longest((repeat(0, x) for x in zeros), (repeat(1, x) for x in ones), fillvalue=[]))))

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))
def chepner_2(zeros, ones):
    zero_groups = (repeat(0, x) for x in zeros)
    one_groups = (repeat(1, x) for x in ones)    
    return list(chain.from_iterable(roundrobin(zero_groups, one_groups)))
def chepner_2__superb_rain(zeros, ones):
    return list(chain.from_iterable(map(repeat, cycle([0, 1]), roundrobin(zeros, ones))))

def Alain_T(zeros, ones):
    return [B for N,P in zip(zeros+[0],ones+[0]) for B in ([0]*N+[1]*P)]

def python_user(zeros, ones):
    res = [None] * (len(ones) + len(zeros))

    res[::2] = ([0]*n for n in zeros)
    res[1::2] = ([1]*n for n in ones)

    res = [y for x in res for y in x]
    return res

def superb_rain(zeros, ones):
    zeros = iter(zeros)
    output = [0] * next(zeros)
    for o in ones:
        output += (1,) * o
        output += (0,) * next(zeros)
    return output

def superb_rain2(zeros, ones):
    z = iter(zeros).__next__
    output = [0] * z()
    for o in ones:
        output += (1,) * o
        output += (0,) * z()
    return output

funcs = [
    baseline,
    original,
    chepner_1,
    chepner_2,
    chepner_2__superb_rain,
    Alain_T,
    python_user,
    superb_rain,
    superb_rain2,
    ]

zeros = [2,3,1,2]
ones = [3,4,5]
number = 10**5

expect = original(zeros, ones)
for func in funcs:
    print(func(zeros, ones) == expect, func.__name__)
print()

tss = [[] for _ in funcs]
for _ in range(4):
    for func, ts in zip(funcs, tss):
        t = min(timeit.repeat(lambda: func(zeros, ones), number=number)) / number
        ts.append(t)
        print(*('%.2f us ' % (1e6 * t) for t in ts[1:]), func.__name__)
    print()