在 numpy arrays 中添加项目，其中每个项目都有一个关联的“类型”

Question

I'm wondering if it's possible to add numpy arrays together item-wise, where each item has a "type" associated with it.

For example, say I have the following 2 arrays:

array1 = [[2 'x'] [3 'y'] [4 'z']]  
array2 = [[10 'x'] [6 'z']]

And I would like to produce this array:

array3 = [[12 'x'] [3 'y'] [10 'z']]

Is there a straightforward way to do this using numpy ? Thanks!

Answer 1

You can either use numpy directly or make advantage of pandas built-in groupby method.

Numpy Implementation

import numpy as np

a1 = np.array([[2, 'x'], [3, 'y'], [4 ,'z']])
a2 = np.array([[10, 'x'] ,[6 ,'z']])

# Stack the two arrays together
a = np.vstack([a1,a2])

# Define "groups"
grp = np.unique(a[:,1])

# Groupby each group in grp and sum the other column
np.array([[a[a[:,1]==g][:,0].astype(float).sum(), g]  for g in grp])
array([['12.0', 'x'],
       ['3.0', 'y'],
       ['10.0', 'z']], dtype='<U32')

Pandas Implementation

import pandas as pd

# Convert a into a pandas DataFrame
df = pd.DataFrame(a, columns=list('ab'))
# Cast first column type
df['a'] = df['a'].astype(int)

# Use groupby sum
df.groupby('b').sum().reset_index()[['a','b']].to_numpy()
array([[12, 'x'],
       [3, 'y'],
       [10, 'z']], dtype=object)