Rearanging numpy array

Question

I need to parse the time-dependent positions of two objects and I get data in numpy array:

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

so that the first column represents a position, the second one the time point A was in that particular position and the last column time at which point B was in that position. It is guaranteed that data is consistent, that is if any two rows have the same times - they have the same position, in pseudocode:

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

What I would love to have is to somehow recast this array so it would enumerate all available times and according positions, such as:

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

Here, first column is the time, second is position of point A and third one is position of point B. np.nan should be assigned when I do not have the information on position of a point. What I currently do is rip the data array into two separate arrays:

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

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

and the I re-merge the as done in John Galt's answer here. This works somehow, but I am really hoping there may be something of a better solution. Can anyone kick me in the right direction?

Answer 1

Here's one approach using NumPy array initialization and assignment -

# 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 acts like a godsend here, as it gives us the locations where we need to put in a and b from data in the already sorted array all_idx and is known to be quite efficient.

Output for given sample data -

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