Idr
Idr

Reputation: 6250

Sort on multiple NumPy arrays

I am creating a 2 dimensional numpy array that contains stock returns. I want to sum the return every 2 days, and if the sum is in the top two, I will set every element in a similar shaped array to True.

For example, returns below is the daily returns for four different stocks.

returns=np.array([
[0, 0, 4, 8],
[7, 5, 4, 1],
[10, 5, 7, 6],
[7, 5, 4, 2]])

For the first two days, columns 2 and 3 (using 0 based indexing) have the highest sums. For the second set of two days, columns 0 and 2 have the highest sums. The output array I want is

bools=np.array([
[False, False, True, True],
[False, False, True, True],
[True, False, True, False],
[True, False, True, False]])

What are good ways to accomplish this?

If there are ties with the sums of two days, I want to use another similarly shaped numpy array as tiebreakers.

For example, if

returns=np.array([
[0, 9, 4, 8],
[7, 5, 4, 0],
[10, 5, 7, 6],
[7, 5, 4, 2]])

For the first set of two days, columns 2 and 3 are tied for the second highest sum. I want to decide the tiebreaker by the greatest value in the last row for the tied columns so that the tie break between columns 2 and 3 look at tiebreaks[1][2] vs tiebreaks[1][3] (4 v 5), and that the ultimate output is bools2.

tiebreaks=np.array([
[0, 0, 1, 1],
[2, 3, 4, 5],
[0, 5, 7, 6],
[-7, 5, -4, 2]])

bools2=np.array([
[False, True, False, True],
[False, True, False, True],
[True, False, True, False],
[True, False, True, False]])

Thanks for your help.

Upvotes: 2

Views: 1437

Answers (1)

Sven Marnach
Sven Marnach

Reputation: 602595

You can use numpy.lexsort() to get the indices that sort your arrays using prices as primary key and names as secondary key. Applying advanced indexing using these indices yields the sorted arrays:

col_indices = numpy.lexsort((names, prices))
row_indices = numpy.arange(len(names))[:, None]
print(prices[row_indices, col_indices])
print(names[row_indices, col_indices])

(Note that in your example, names and prices don't have compatible shapes.)

Upvotes: 1

Related Questions