javadl4jnd4j

How to iterate over indices of a matrix?


In Python when we want to iterate over a matrix with an arbitrary dimension, we can use this line of code:

for index in np.ndindex(data.shape[2:]):

for example :

> for index in np.ndindex(3, 2, 1):
>     print(index) (0, 0, 0) (0, 1, 0) (1, 0, 0) (1, 1, 0) (2, 0, 0) (2, 1, 0)

in java, in an easy way, we can do it with a determined number of for loop, but the prerequisite is knowledge about the dimension. But in arbitrary dimensions, the algorithm must be more complicated.

Is there any built-in method in ND4J lib for iterating over indices?


Solution

  • In nd4j, we have a NDIndexIterator that allows you to iterate over the coordinates.

    Here is the example:

    NdIndexIterator shapeIter = new NdIndexIterator(2, 2);
    //import org.nd4j.linalg.api.iter.NdIndexIterator;
    
    long[][]  possibleSolutions = new long[][] {{0, 0}, {0, 1}, {1, 0}, {1, 1},};
    for (int i = 0; i < 4; i++) {
        assertArrayEquals(possibleSolutions[i], shapeIter.next());
    }