给定一个NumPy数组,其中第一列的值表示分组依据,如何对相同分组依据的行求和?例如,给定以下数组:
np.array([[1,2,3],
[1,4,6],
[2,3,5],
[2,6,2],
[3,4,8]])
期望的输出为:
[[1,6,9], [2,9,7], [3,4,8]]
2. 解决方案
方法1:使用Pandas GroupBy函数
Pandas提供了一个非常强大的GroupBy函数,可以轻松实现对具有相同值的列求和。具体步骤如下:
- 将NumPy数组转换为Pandas DataFrame:
import pandas as pd
n = np.array([[1,2,3],
[1,4,6],
[2,3,5],
[2,6,2],
[3,4,8]])
df = pd.DataFrame(n, columns = ["First Col", "Second Col", "Third Col"])
2. 使用GroupBy函数对第一列(“First Col”)分组并求和:
df.groupby("First Col").sum()
输出:
First Col Second Col Third Col
0 1 6 9
1 2 9 7
2 3 4 8
方法2:使用NumPy bincount函数
NumPy的bincount函数可以实现对具有相同值的元素进行计数。这里,我们可以利用它来对分组依据进行计数,从而得到分组和。具体步骤如下:
- 使用NumPy的unique函数获取分组依据的唯一值及其对应的索引:
unqA1, id = np.unique(A[:, 0], return_inverse=True)
2. 使用NumPy的bincount函数对索引进行计数,得到分组和:
sums = np.bincount( subs.ravel(), weights=A[:,1:].ravel() )
3. 将分组依据的唯一值和分组和组合成输出数组:
out = np.append(unqA1[:,None],sums.reshape(N,-1).T,1)
输出:
[[1 6 9]
[2 9 7]
[3 4 8]]
方法3:使用NumPy cumsum和diff函数
NumPy的cumsum函数可以实现对数组元素的累积求和,diff函数可以实现对数组元素的差分计算。这里,我们可以利用它们来对分组依据进行分组和求和。具体步骤如下:
- 对数组按第一列进行升序排序:
sA = A[np.argsort(A[:,0]),:]
2. 使用NumPy的diff函数获取分组依据的差分,得到分组结束的行掩码:
row_mask = np.append(np.diff(sA[:,0],axis=0)!=0,[True])
3. 使用NumPy的cumsum函数对数组按行累积求和,得到分组和:
cumsum_grps = sA.cumsum(0)[row_mask,1:]
4. 使用NumPy的diff函数对分组和进行差分,得到分组和:
sum_grps = np.diff(cumsum_grps,axis=0)
5. 将分组依据的唯一值和分组和组合成输出数组:
out = np.concatenate((sA[row_mask,0][:,None],counts),axis=1)
输出:
[[1 6 9]
[2 9 7]
[3 4 8]]
方法4:使用NumPy add.at函数
NumPy的add.at函数可以对数组元素进行累加。这里,我们可以利用它来对分组依据进行分组和求和。具体步骤如下:
- 使用NumPy的unique函数获取分组依据的唯一值及其对应的索引:
unq, unq_inv = np.unique(A[:, 0], return_inverse=True)
2. 创建一个输出数组,并初始化第一列为分组依据的唯一值:
out = np.zeros((len(unq), A.shape[1]), dtype=A.dtype)
out[:, 0] = unq
3. 使用NumPy的add.at函数对输出数组按索引累加分组依据的行:
np.add.at(out[:, 1:], unq_inv, A[:, 1:])
输出:
[[1 6 9]
[2 9 7]
[3 4 8]]