Challenge - Big Fish

10 fish occupy a 5x5x5 grid of water 🐟. Each fish decides to move to a new (i,j,k) location given by the 2-d array below. If multiple fish end up occupying the same cell, the biggest fish eats the smaller fish. Determine which fish survive.

import numpy as np

locs = np.array([
    [0,0,0],
    [1,1,2],
    [0,0,0],
    [2,1,3],
    [5,5,4],
    [5,0,0],
    [5,0,0],
    [0,0,0],
    [2,1,3],
    [1,3,1]
])

generator = np.random.default_rng(1010)
weights = generator.normal(size=10)

print(weights)

[-1.69870017  0.53799701 -0.22561399 -1.09020894  0.55391264 -1.50115445
  0.44545933  1.3448172  -1.12364327  0.21216015]

# Calculate the weights in descending order
reverse_sorted_weights = np.argsort(weights)[::-1]

# Get the first (since its heaviest to lightest) unique coordinates to see which one survive
_, indices = np.unique(locs[reverse_sorted_weights], axis=0, return_index=True)

# Get the fish that survive
surive_weight_indices = reverse_sorted_weights[indices]
surive_weight_indices, weights[surive_weight_indices]

(array([7, 1, 9, 3, 6, 4]),
 array([ 1.3448172 ,  0.53799701,  0.21216015, -1.09020894,  0.44545933,
         0.55391264]))