python 使用“模板列表”創建新列表

Question

假設我有：

x1 = [1, 3, 2, 4]

和：

x2 = [0, 1, 1, 0]

具有相同的形狀

現在我想“將 x2 放在 x1 上”並總結與 x2 的數量相對應的所有 x1 的數量

所以最終結果是：

end = [1+4 ,3+2]  # end[0] is the sum of all numbers of x1 where a 0 was in x2

這是一個使用列表來進一步澄清問題的幼稚實現

store_0 = 0
store_1 = 0
x1 = [1, 3, 4, 2]
x2 = [0, 1, 1, 0]
for value_x1 ,value_x2 in zip(x1 ,x2):
    if value_x2 == 0:
        store_0 += value_x1
    elif value_x2 == 1:
        store_1 += value_x1

所以我的問題是：有沒有一種方法可以在 numpy 中實現這一點，而不使用循環或通常更快？

Answer 1

在這個特定的示例中（通常，對於unique 、 duplicated和groupby類型的操作）， pandas比純numpy解決方案更快：

使用Series的pandas方式（信用：與@mcsoini 的答案非常相似）：

def pd_group_sum(x1, x2):
    return pd.Series(x1, index=x2).groupby(x2).sum()

一個純粹的numpy方式，使用np.unique和一些花哨的索引：

def np_group_sum(a, groups):
    _, ix, rix = np.unique(groups, return_index=True, return_inverse=True)
    return np.where(np.arange(len(ix))[:, None] == rix, a, 0).sum(axis=1)

注意：更好的純numpy方式的靈感來自@Woodford 的回答：

def selsum(a, g, e):
    return a[g==e].sum()

vselsum = np.vectorize(selsum, signature='(n),(n),()->()')

def np_group_sum2(a, groups):
    return vselsum(a, groups, np.unique(groups))

另一種純粹的numpy方式的靈感來自@mapf 關於使用argsort()的評論。 這本身已經花費了 45 毫秒，但我們可以嘗試基於np.argpartition(x2, len(x2)-1)的東西，因為在下面的基准測試中它本身只需要 7.5 毫秒：

def np_group_sum3(a, groups):
    ix = np.argpartition(groups, len(groups)-1)
    ends = np.nonzero(np.diff(np.r_[groups[ix], groups.max() + 1]))[0]
    return np.diff(np.r_[0, a[ix].cumsum()[ends]])

（稍作修改）示例

x1 = np.array([1, 3, 2, 4, 8])  # I added a group for sake of generality
x2 = np.array([0, 1, 1, 0, 7])

>>> pd_group_sum(x1, x2)
0    5
1    5
7    8

>>> np_group_sum(x1, x2)  # and all the np_group_sum() variants
array([5, 5, 8])

速度

n = 1_000_000
x1 = np.random.randint(0, 20, n)
x2 = np.random.randint(0, 20, n)

%timeit pd_group_sum(x1, x2)
# 13.9 ms ± 65.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit np_group_sum(x1, x2)
# 171 ms ± 129 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np_group_sum2(x1, x2)
# 66.7 ms ± 19.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np_group_sum3(x1, x2)
# 25.6 ms ± 41.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

通過 pandas 更快，部分原因是numpy 問題 11136 。

Answer 2

>>> x1 = np.array([1, 3, 2, 7])
>>> x2 = np.array([0, 1, 1, 0])
>>> for index in np.unique(x2):
>>>     print(f'{index}: {x1[x2==index].sum()}')
0: 8
1: 5
>>> # or in one line
>>> [(index, x1[x2==index].sum()) for index in np.unique(x2)]
[(0, 8), (1, 5)]

Answer 3

pandas 單線可以嗎？

store_0, store_1 = pd.DataFrame({"x1": x1, "x2": x2}).groupby("x2").x1.sum()

或者作為字典，對於x2中的任意多個值：

pd.DataFrame({"x1": x1, "x2": x2}).groupby("x2").x1.sum().to_dict()

Output：

{0: 5, 1: 5}

Answer 4

使用壓縮

from itertools import compress
result = [sum(compress(x1,x2)),sum(compress(x1, (map(lambda x: not x,x2))))]

Answer 5

這會將您的循環擴展到更多的值。 我想不出一個 numpy 單線來做到這一點。

sums = [0] * 10000
for vx1,vx2 in zip(x1,x2):
    sums[vx2] += vx1

Answer 6

通過將第二個列表轉換為 Boolean 數組，您可以使用它來索引第一個：

import numpy as np

x1 = np.array([1, 3, 2, 4])
x2 = np.array([0, 1, 1, 0], dtype=bool)

end = [np.sum(x1[~x2]), np.sum(x1[x2])]
end

[5, 5]

編輯：如果x2的值可以大於 1，則可以使用列表推導：

x1 = np.array([1, 3, 2, 4])
x2 = np.array([0, 1, 1, 0])

end = [np.sum(x1[x2 == i]) for i in range(max(x2) + 1)]

Answer 7

這擴展了 Tim Roberts 在開始時建議的解決方案，但將說明X2具有多個值，即非二進制。 這里這些值是嚴格相鄰的，因為 for 循環使用rng的range ，但它可以擴展，以便 x2 具有不相鄰的值，例如 [0 2 2 2 1 4] <- no 3's 而用於此示例的randint將返回一個類似於 [0 1 1 3 4 2] 的向量。

import numpy as np
rng = 5 # Range of values for x2 i.e [0 1 2 3 4]
x1 = np.random.randint(20, size=10000) #random vector of size 10k
x2 = np.random.randint(5, size=10000) # inexing vector size 10k with range (0-4)


store = []
for i in range(rng): # loop and append to list
    store.append(x1[x2==i].sum())