NumPy将函数应用于与另一个numpy数组相对应的行组

Question

我有一个NumPy数组，每一行都表示一些（x，y，z）坐标，如下所示：

a = array([[0, 0, 1],
       [1, 1, 2],
       [4, 5, 1],
       [4, 5, 2]])

我还有另一个NumPy数组，该数组的z坐标具有唯一的值，如下所示：

b = array([1, 2])

我如何将一个函数（我们称其为“ f”）应用于与b中的值相对应的a中的每组行？例如，b的第一个值为1，因此我将得到a的所有行，这些行的z坐标为1。然后，我将一个函数应用于所有这些值。

最后，输出将是与b形状相同的数组。

我正在尝试将其向量化，以使其尽可能快。谢谢！

预期输出的示例（假设f为count（））：

c = array([2, 2])

因为在数组a中有2行在数组b中的z值为1，在数组a中又有2行在数组b中的z值为2。

一个简单的解决方案是像这样遍历数组b：

for val in b:
    apply function to a based on val
    append to an array c

我的尝试：

我尝试做这样的事情，但是它只是返回一个空数组。

func(a[a[:, 2]==b])

Answer 1

问题是，具有相同Z的行的组可以具有不同的大小，因此您不能将它们堆叠到一个3D numpy数组中，这将允许轻松地沿第三个维度应用函数。一种解决方案是使用for循环，另一种解决方案是使用np.split：

a = np.array([[0, 0, 1],
              [1, 1, 2],
              [4, 5, 1],
              [4, 5, 2],
              [4, 3, 1]])


a_sorted = a[a[:,2].argsort()]

inds = np.unique(a_sorted[:,2], return_index=True)[1]

a_split = np.split(a_sorted, inds)[1:]

# [array([[0, 0, 1],
#         [4, 5, 1],
#         [4, 3, 1]]),

#  array([[1, 1, 2],
#         [4, 5, 2]])]

f = np.sum

result = list(map(f, a_split))
# [19, 15]

但是恕我直言，最好的解决方案是使用FBruzzesi建议的pandas和groupby。然后，您可以将结果转换为numpy数组。

EDIT：为完整性起见，这是另外两个解决方案

列表理解：

b = np.unique(a[:,2])
result = [f(a[a[:,2] == z]) for z in b]

熊猫：

df = pd.DataFrame(a, columns=list('XYZ'))
result = df.groupby(['Z']).apply(lambda x: f(x.values)).tolist()

这是我为a = np.random.randint(0, 100, (n, 3))获得的性能图：

您可以看到，大约n = 10^5的“拆分解决方案”是最快的，但此后熊猫解决方案的性能更好。

Answer 2

如果允许您使用熊猫：

import pandas as pd
df=pd.DataFrame(a, columns=['x','y','z'])

df.groupby('z').agg(f)

这里f可以是任何处理分组数据的自定义函数。

数值示例：

a = np.array([[0, 0, 1],
              [1, 1, 2],
              [4, 5, 1],
              [4, 5, 2]])
df=pd.DataFrame(a, columns=['x','y','z'])
df.groupby('z').size()

z
1    2
2    2
dtype: int64

注意.size是计算每组行数的方法。

为了使其成为纯粹的numpy，也许这可以适合您的情况：

tmp = np.array([a[a[:,2]==i] for i in b])
tmp 
array([[[0, 0, 1],
        [4, 5, 1]],

       [[1, 1, 2],
        [4, 5, 2]]])

这是每个数组组成的数组。

Answer 3

c = np.array([])
for x in np.nditer(b):
    c = np.append(c, np.where((a[:,2] == x))[0].shape[0])

输出：

[2. 2.]

Answer 4

我投票支持Andreas K的解决方案。类似的问题是asked here，该解决方案与答案之间的解决方案相同。要注意的最重要的事情是将解决方案分为三个子类别：基于pandas的，基于numpy的和基于numpy_indexed的子类别。 df.groupby方法对我来说似乎是最简单的方法，基于numpy的分组更为复杂，因为它需要了解np.argsort，np.unique和np.split方法的组合。但是，仅使用numpy操作合并这三种方法应该仍然是进行分组的主要方法。