python - 有效地计算 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 数组中不会出现超过一次。