重新排列numpy数组

Question

我需要解析两个物体的时间相关位置，并且它们的数据以numpy数组的形式给出：

data = [[0, 1, 2],
        [1, 4, 3],
        [2, 2, 1]]

第一列代表一个位置，第二列代表时间点A在该特定位置的时间，最后一列代表时间点B在该位置的时间。可以保证数据一致性，也就是说，如果任意两行具有相同的时间，则它们具有相同的位置，在伪代码中：

data[row1,1] == data[row2,1]  <=>  data[row1,0] == data[row2,0]
data[row1,2] == data[row2,2]  <=>  data[row1,0] == data[row2,0]

我希望的是以某种方式重新排列这个数组，使其枚举所有可用的时间和相应的位置，例如：

parsed = [[1, 0, 2],
          [2, 2, 0],
          [3, np.nan, 1],
          [4, 1, np.nan]]

这里，第一列是时间，第二列是点A的位置，第三列是点B的位置。当我没有关于某个点的位置信息时，应该分配np.nan。目前我的做法是将数据数组分成两个单独的数组：

    moments = set (data [:, 1:3].flatten())

    for each in moments:
        a = data[:,[1,0]][pos[:,1] == each]
        b = data[:,[2,0]][pos[:,2] == each]

我按照John Galt在这里的答案所述的方法重新合并了它们。这种方法可以运行，但我真的希望有更好的解决方案。有人能否指点一下我正确的方向？

- Stipe Galić

---

只是为了将来参考的注释；我在问题中表述得有些笨拙。当我写道“数据一致是有保证的，也就是说，如果任意两行具有相同的时间，则它们具有相同的位置”时，这是正确的，但是反过来则不是，正如在“伪代码”中使用符号“<==>”所建议的那样——一个点具有相同的位置和不同的时间是完全可以的。 - Stipe Galić

Answer 1

这里有一种使用NumPy数组初始化和赋值的方法 -

# Gather a and b indices. Get their union, that represents all posssible indices
a_idx = data[:,1]
b_idx = data[:,2]
all_idx = np.union1d(a_idx, b_idx)

# Setup o/p array 
out = np.full((all_idx.size,3),np.nan)

# Assign all indices to first col
out[:,0] = all_idx

# Determine the positions of a indices in all indices and assign first col data
out[np.searchsorted(all_idx, a_idx),1] = data[:,0]
# Similarly for b
out[np.searchsorted(all_idx, b_idx),2] = data[:,0]

np.searchsorted在这里就像天赐良机一样，因为它给出了我们需要将a和b从data放入已排序数组all_idx中的位置，并且已知非常高效。

给定示例数据的输出 -

In [104]: out
Out[104]: 
array([[  1.,   0.,   2.],
       [  2.,   2.,   0.],
       [  3.,  nan,   1.],
       [  4.,   1.,  nan]])

- Divakar

---

天啊，这个东西飞快！所以，你使用union1d而不是我用set()的那个东西。你用nans初始化相应的最终数组，并使用searchsorted()根据时刻向量将其复制到相应的位置，缺乏更好的词语？我会点赞并且如果测试数据没问题，我会接受它作为答案。 - Stipe Galić

根据所有可能的索引顺序排序，这让我们能够利用超级高效的“searchsorted”。在大多数情况下，NumPy都能如此高效。@StipeGalić - Divakar

Answer 2

在没有更好的想法的情况下，让我来一个Pandas的一行代码解决方案。免责声明：它的运行速度比Divakar的纯Numpy解决方案慢100倍：

df = pd.DataFrame(data)
pd.concat([df.set_index(ix)[0] for ix in [1,2]], axis=1).reset_index().values
#array([[  1.,   0.,   2.],
#       [  2.,   2.,   0.],
#       [  3.,  nan,   1.],
#       [  4.,   1.,  nan]])

- DYZ

---

感谢您的努力。我实际上在使用Pandas，并且有10MB以上的数据，我很感激这两个数量级的加速。;) - Stipe Galić