python - 递增地增加列表中的非零元素

Question

如果我得到这个列表

a = [1,0,0,1,0,0,0,1]

我想把它变成

a = [1,0,0,2,0,0,0,3]

Answer 1

解决方案 #1 和 #2 的设置

from itertools import count

to_add = count()
a = [1,0,0,1,0,0,0,1]

解决方案#1

>>> [x + next(to_add) if x else x for x in a]
[1, 0, 0, 2, 0, 0, 0, 3]

解决方案 #2，hacky 但有趣

>>> [x and x + next(to_add) for x in a]
[1, 0, 0, 2, 0, 0, 0, 3]

解决方案 #3 和 #4 的设置

import numpy as np
a = np.array([1,0,0,1,0,0,0,1])

解决方案#3

>>> np.where(a == 0, 0, a.cumsum())
array([1, 0, 0, 2, 0, 0, 0, 3])

解决方案#4(我最喜欢的一个)

>>> a*a.cumsum()
array([1, 0, 0, 2, 0, 0, 0, 3])

所有cumsum 解决方案都假设a 的非零元素都是1。

---

时间:

# setup
>>> a = [1, 0, 0, 1, 0, 0, 0, 1]*1000
>>> arr = np.array(a)
>>> to_add1, to_add2 = count(), count()
# IPython timings @ i5-6200U CPU @ 2.30GHz (though only relative times are of interest)
>>> %timeit [x + next(to_add1) if x else x for x in a] # solution 1
669 µs ± 3.59 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit [x and x + next(to_add2) for x in a] # solution 2
673 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit np.where(arr == 0, 0, arr.cumsum()) # solution 3
34.7 µs ± 94.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> %timeit arr = np.array(a); np.where(arr == 0, 0, arr.cumsum()) # solution 3 with array creation
474 µs ± 14.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
>>> %timeit arr*arr.cumsum() # solution 4
23.6 µs ± 131 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> %timeit arr = np.array(a); arr*arr.cumsum() # solution 4 with array creation
465 µs ± 6.82 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)