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, которые имеют 1 в координате z. Затем я применяю функцию ко всем этим значениям.

В конце концов, вывод будет массив той же формы, что и b.

Я пытаюсь векторизовать это, чтобы сделать его как как можно быстрее. Спасибо!

Пример ожидаемого вывода (при условии, что f равно count ()):

c = array([2, 2])

, поскольку в массиве a есть 2 строки, которые имеют значение az 1 в массиве b и также 2 строки в массиве a, которые имеют значение az 2 в массиве b.

Тривиальным решением было бы перебрать массив 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-l oop, другое - 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]

Но imho лучшее решение - использовать pandas и groupby, как это было предложено FBruzzesi. , Затем вы можете преобразовать результат в массив numpy.

EDIT : для полноты рассмотрим два других решения

Понимание списка:

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

Pandas:

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 «разделенное решение» является самым быстрым, но после этого решение pandas работает лучше.

Answer 2

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

Выход:

[2. 2.]

Answer 3

Если вам разрешено использовать pandas:

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

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

Здесь f может быть любой пользовательской функцией, работающей с сгруппированными данными.

Numeri c пример:

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

, который является массивом с каждой группой массивов.