Find equivalent of "enumerate" of NumPy arrays

In Python, the enumerate function is a convenient way to iterate over a sequence while keeping track of the index of each item. However, when working with NumPy arrays, the built-in enumerate function does not work as it only supports Python sequences. Fortunately, NumPy provides an alternative functions, which are specifically designed to iterate over multi-dimensional arrays while also returning the corresponding index of each element. Let discuss some methods below:

1. Using "np.ndenumerate()" :

You can used np.ndenumerate() method to find iterate NumPy arrays. np.ndenumerate() returns an iterator of index-value pairs for a given array. Here’s an example given below:


  • The code creates a 2D NumPy array with the values [[1, 2], [3, 4]].

  • It then uses the np.ndenumerate() function to iterate over the indices and values of the array.

  • The np.ndenumerate() function returns an iterator that yields the index and corresponding value of each element in the array.

  • In this case, the loop iterates over each element in the array and prints its index and value to the console.

In terms of memory, the code creates a 2D NumPy array of size 2x2, which requires 16 bytes of memory. The memory required for the iterator created by np.ndenumerate() is negligible compared to the memory used by the array. np.ndenumerate() is more useful because it returns both the indices and the corresponding values.

2. Using "np.ndindex()" :

Another option is to utilize the np.ndindex() function, which generates an iterator of index tuples for a specified array shape. These index tuples can then be utilized to access the corresponding values within the array. An example to illustrate this is provided below:


  • The above code first creates a 2D array with the values [[1,2], [3,4]].

  • Then, it uses the np.ndindex() function to create an iterator over the indices of the array. This iterator returns a tuple of indices for each iteration.

  • The for loop iterates over this iterator and extracts the value of the array at the current index using the array[index] syntax.

  • The index and value are then printed using the print() function.

In terms of memory, the code creates a 2D NumPy array of size 2x2, which requires 16 bytes of memory. The np.ndindex() function is faster and more memory-efficient because np.ndindex() returns only the indices of the array.