有效地计算 numpy 列中的重复值并附加计数

Question

我有一个表示有向图的数据集。第一列是源节点，第二列是目标节点，我们可以忽略第三列（本质上是一个权重）。比如：

0 1 3
0 13 1
0 37 1
0 51 1
0 438481 1
1 0 3
1 4 354
1 10 2602
1 11 2689
1 12 1
1 18 345
1 19 311
1 23 1
1 24 366
...

我想做的是为每个节点附加出度。例如，如果我只是添加节点 0 的出度，我会：

0 1 3 5
0 13 1 5
0 37 1 5
0 51 1 5
0 438481 1 5
1 0 3
...

我有一些代码可以做到这一点，但它非常慢，因为我使用的是 for 循环：

import numpy as np

def save_degrees(X):
    new_col = np.zeros(X.shape[0], dtype=np.int)
    X = np.column_stack((X, new_col))
    node_ids, degrees = np.unique(X[:, 0], return_counts=True)
    # This is the slow part.
    for node_id, deg in zip(node_ids, degrees):
        indices = X[:, 0] == node_id
        X[:, -1][indices] = deg
    return X

train_X = np.load('data/train_X.npy')
train_X = save_degrees(train_X)
np.save('data/train_X_degrees.npy', train_X)

有没有更有效的方法来构建这种数据结构？

Answer 1

您可以使用numpy.unique。

假设你的输入数据在数组data：

In [245]: data
Out[245]: 
array([[     0,      1,      3],
       [     0,     13,      1],
       [     0,     37,      1],
       [     0,     51,      1],
       [     0, 438481,      1],
       [     1,      0,      3],
       [     1,      4,    354],
       [     1,     10,   2602],
       [     1,     11,   2689],
       [     1,     12,      1],
       [     1,     18,    345],
       [     1,     19,    311],
       [     1,     23,      1],
       [     1,     24,    366],
       [     2,     10,      1],
       [     2,     13,      3],
       [     2,     99,      5],
       [     3,     25,     13],
       [     3,     99,     15]])

在第一列中查找唯一值，以及“逆”数组和每个唯一值的出现次数：

In [246]: nodes, inv, counts = np.unique(data[:,0], return_inverse=True, return_counts=True)

您的出度列是counts[inv]：

In [247]: out_degrees = counts[inv]

In [248]: out_degrees
Out[248]: array([5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9, 9, 3, 3, 3, 2, 2])

这假设一对 (source_node, target_node) 在 data 数组中不会出现多次。

Answer 2

np.unique 确实在这里做得很好，正如其他一些答案中所解释的那样。

你可能还想看看numpy_indexed（免责声明：我是它的作者）；它可以以相同的效率做同样的事情，但也支持许多其他功能，这在处理图形时往往非常有用；或一般稀疏/锯齿状的数据结构。

它还专门针对您的问题提供了一个干净的单行解决方案：

import numpy_indexed as npi
X = np.column_stack((X, npi.multiplicity(X[:, 0])))

Answer 3

你可以试试这个，通常X[:, 0] == node_id 当你有很多不同的节点时会很耗时。您可以先按第一列对数据进行排序，然后通过重复计数从中创建一个新的计数列：

train_X = train_X[train_X[:, 0].argsort()]
_, counts = np.unique(train_X[:,0], return_counts=True)
np.hstack((train_X, np.repeat(counts, counts)[:, None]))

# array([[     0,      1,      3,      5],
#        [     0,     13,      1,      5],
#        [     0,     37,      1,      5],
#        [     0,     51,      1,      5],
#        [     0, 438481,      1,      5],
#        [     1,      0,      3,      9],
#        [     1,      4,    354,      9],
#        [     1,     10,   2602,      9],
#        [     1,     11,   2689,      9],
#        [     1,     12,      1,      9],
#        [     1,     18,    345,      9],
#        [     1,     19,    311,      9],
#        [     1,     23,      1,      9],
#        [     1,     24,    366,      9]])

或者你可以使用 pandas groupby：

import pandas as pd
pd.DataFrame(train_X).pipe(lambda x: x.assign(size = x.groupby([0])[0].transform('size'))).values

#array([[     0,      1,      3,      5],
#       [     0,     13,      1,      5],
#       [     0,     37,      1,      5],
#       [     0,     51,      1,      5],
#       [     0, 438481,      1,      5],
#       [     1,      0,      3,      9],
#       [     1,      4,    354,      9],
#       [     1,     10,   2602,      9],
#       [     1,     11,   2689,      9],
#       [     1,     12,      1,      9],
#       [     1,     18,    345,      9],
#       [     1,     19,    311,      9],
#       [     1,     23,      1,      9],
#       [     1,     24,    366,      9]])

Answer 4

这是一种注重性能的矢量化方法 -

def argsort_unique(idx):
    # Original idea : http://stackoverflow.com/a/41242285/3293881 
    n = idx.size
    sidx = np.empty(n,dtype=int)
    sidx[idx] = np.arange(n)
    return sidx

def count_and_append(a): # For sorted arrays
    a0 = a[:,0]
    sf0 = np.flatnonzero(a0[1:] != a0[:-1])+1
    shift_idx = np.concatenate(( [0] , sf0, [a0.size] ))
    c = shift_idx[1:] - shift_idx[:-1]
    out_col = np.repeat(c,c)
    return np.column_stack((a, out_col))

def count_and_append_generic(a): # For generic (not necessarily sorted) arrays
    sidx = a[:,0].argsort()
    b = a[sidx]
    return count_and_append(b)[argsort_unique(sidx)]

示例运行 -

In [70]: a # Not sorted case
Out[70]: 
array([[     1,     18,    345],
       [     1,     23,      1],
       [     0,     13,      1],
       [     0,     37,      1],
       [     2,     99,      5],
       [     0,      1,      3],
       [     2,     13,      3],
       [     1,      4,    354],
       [     1,     24,    366],
       [     0, 438481,      1],
       [     1,     12,      1],
       [     1,     11,   2689],
       [     1,     19,    311],
       [     2,     10,      1],
       [     3,     99,     15],
       [     0,     51,      1],
       [     3,     25,     13],
       [     1,      0,      3],
       [     1,     10,   2602]])

In [71]: np.allclose(count_and_append_generic(a), save_degrees(a))
Out[71]: True

如果输入数组已经按第一列排序，只需使用count_and_append(a)。