在 Python 中生成和 zip 两个列表的最干净有效的方法

Question

Given these two lists

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

(with the condition that always len(zeros) == len(ones) + 1 )

I want to create one list with alternating 0's and 1's of the magnitude mentioned in the list. I can achieve this by:

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]

Is this however the most efficient / clean way? I get a performance of 2.66 µs ± 78.8 ns, but I still have the idea this could be done in a one-liner, and maybe more efficient

Answer 1

Two apparently faster solutions, using the "trick" to special-treat the first zero rather than the last, using it to initialize output .

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

(As @schwobaseggl pointed out, tuples made it about 30% faster.)

Benchmark results:

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

Benchmark code:

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()

Answer 2

Zip with a list comprehension should do the trick.

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

output = [B for N,P in zip(zeros,ones+[0]) for B in [0]*N+[1]*P]

print(output)

[0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 0]

Note the ones+[0] is to ensure that you don't drop the last value from the zeroes list in the zip operation.

Answer 3

You can use itertools.chain , zip_longest , and itertools.repeat to create a not-too-confusing one-liner.

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

On my machine, this took 3.34µs. More importantly, it's the wrapper to list that takes this time. The iterator itself produces the elements on demand, if you don't actually need them all at once.

---

list(chain.from_iterable(
        chain.from_iterable(zip_longest((repeat(0, x) for x in zeros),
                                        (repeat(1, x) for x in ones),
                                        fill_value=[]))))

• (repeat(0, x) for x in zeros) creates a sequence of repeat objects that represent your runs of 0s; likewise to create the groups of 1s.
• zip_longest zips them together into a sequence of pairs, adding a do-nothing empty list to balance the extra value in zeros
• chain.from_iterable flattens that sequence (from (a, b), (c, d) to (a, b, c, d) .
• The outer chain.from_iterable then flattens the repeat objects in to a single sequence, which list turns into a list.

---

You can also simplify the one-liner using the roundrobin recipe from the itertools documentation, which handles both merging the zero groups and one groups, as well as the first round of flattening.

from itertools import cycle, islice, repeat, chain

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

# From the itertools documentation
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))


zero_groups = (repeat(0, x) for x in zeros)
one_groups = (repeat(1, x) for x in ones)
    
print(list(chain.from_iterable(roundrobin(zero_groups, one_groups))))

Answer 4

Given these two lists

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

(with the condition that always len(zeros) == len(ones) + 1 )

I want to create one list with alternating 0's and 1's of the magnitude mentioned in the list. I can achieve this by:

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]

Is this however the most efficient / clean way? I get a performance of 2.66 µs ± 78.8 ns, but I still have the idea this could be done in a one-liner, and maybe more efficient