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”）应用到 a 中与 b 中的值相对应的每个行组？例如，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  # example of a function

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

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

编辑：为了完整起见，这里是另外两个解决方案

列表理解：

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

如果你被允许使用 pandas：

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.]