find_k_largest_elements_in_2d_array.py

import numpy as np

# FIND LARGEST K ELEMENTS IN A 2D ARRAY IN NUMPY

def find_k_largest_elements_in_2d_array(arr: np.array, k=1):
    assert len(arr.shape) == 2
    h, w = arr.shape
    top_indices = np.flip(arr.flatten().argsort()[-k:])
    return np.divmod(top_indices, w)


def test():
    arr = np.random.uniform(size=(1000, 1000))
    k = 10
    top_k = find_k_largest_elements_in_2d_array(arr, k=k)
    for i in range(k):
        assert np.sum(arr[top_k][i] <= arr) == i + 1


if __name__ == '__main__':
    test()