Numpy structured arrays, selecting x,y pair with smallest x-value for several x,y pairs with the same y-value

GeekGeek4 0
Viewed 25

I have a numpy structured array

import numpy as np
arr1 = (np.array([2, 5, 8, 3, 10], dtype=np.int64),np.array([10, 10, 10, 8, 3], dtype=np.int64))

arr1_x = arr1[0]
arr1_y = arr1[1]

arr1_struct = np.empty(arr1_x.shape[0], dtype=[('x', int), ('y', int)])

arr1_struct["x"] = arr1_x
arr1_struct["y"] = arr1_y

As you can see, in the above structured array, I have x and y values. I would need to do the following

  • Check if there are multiple y values that are the same
  • If there are, then I would need to keep only one x,y value pair for that y value.
  • The x value to keep would need to be the smallest x value among the several x values corresponding to the y value in question
  • x,y pairs are unique and not repeated

In the above example, I would need to keep:

  • x = 2 and y = 10 (chosen from the three y=10 values)
  • x = 3 and y = 8
  • x = 10 and y = 3

The following would need to be discarded:

  • x = 5 and y = 10
  • x = 8 and y = 10

Is there a good way to go about this in numpy?

numpy python
1 Answers

First, let's simplify how you create that stuct array a bit, for no other reasons that it's grinding my gears:

data = np.stack(([2, 5, 8, 3, 10], [10, 10, 10, 8, 3]), axis=-1)
data = data.view(np.dtype([('x', int), ('y', int)])).ravel()

After sorting the data and finding the unique values of y, you can apply np.minimum.reduceat to the resulting segments of x:

data = data[np.argsort(data['y'])]
splits = np.r_[0, np.flatnonzero(np.diff(data['y'])) + 1]
x = np.minimum.reduceat(data['x'], splits)
y = data['y'][splits]

Now you can re-create the array the same way as you did before:

result = np.stack((x, y), axis=-1).view(data.dtype).ravel()

This is pretty much how np.unique and pd.DataFrame.groupby work. The nice thing is that this method is insensitive to repeated pairs.

Another way to approach this is using np.lexsort, which can sort both columns simultaneously, so that the first element of each segment is the minimum:

data = data[np.lexsort((data['x'], data['y']))]
splits = np.r_[0, np.flatnonzero(np.diff(data['y'])) + 1]
result = data[splits]

This method is a bit cleaner because it allows you to index directly into the data.

You can also replace the np.diff computations with np.unique using return_index=True. This is slightly more expensive because it sorts the array twice:

data = data[np.lexsort((data['x'], data['y']))]
result = data[np.unique(data['y'], return_index=True)[1]]
Mad Physicist0
Related