Swift algorithm to enumerate a multilinear map, using multiples indexes:[Int]

Question

A multilinear map M has its elements stored in a one-dimension array of length N, with a Shape S defined by S:[Int] = [p,q,r,...] so that q*p*r*... = N. The Shape is of variable size, not known at compile time.

The issue I'm trying to solve is a generic approach to accessing the map's elements using an array of integers, which individual values are coordinates in the Shape S, ex: M[1,3,2], M[2,3,3,3] etc... This is a problem different from a simple enumeration of the map's elements.

One method is to use M[i,j,k] and implement a subscript method. Unfortunately, this approach hardcodes the map's shape, and the algorithm is no longer generic.

Say there's a utility function that returns an element index from a tuple derived from the map's Shape, so that:

func index(_ indexes:[Int]) -> Int {....}

func elementAt(indexes:[Int]) -> Element { 
   return elements_of_the_map[self.index(indexes)]
}

M.elementAt(indexes:[i,j,k]) or M.elementAt(indexes:[i,j,k,l,m]) always work. So the problem at this point is to build the array [i,j,k,...]

Question: Is there an algorithm to efficiently enumerate those indexes? Nested loops won't work since the number of loops isn't known at compile time, and recursive function seem to add a lot of complexity (in particular keeping track of previous indexes).

I'm thinking about an algorithm 'a la' base-x counting, that is adding one unit to the top right index, and moving leftwards one unit if the count exceeds the number of elements by the map's Shape.

Answer 1

Here's the code, it's primitive, but should work. The idea is to increment, right-to-left, to move say to [1,2,2] from [1,2,1] with the shape constraint [2,3,3].

func add_one_unit(shape:[Int],indexes:[Int]) -> [Int]? {
    //Addition is right to left, so we have to reverse the arrays. Shape Arrays are usually very small, so it's fast.
    let uu = Array(indexes.reversed()); //Array to add one index to.
    let shape_reversed = Array(shape.dimensions.reversed()); //Shape array.

    var vv:[Int] = [];
    var move_next:Bool = true;

    for i in 0..<uu.count {
        if move_next {
            if uu[i] < shape_reversed[i] - 1 { //Shape constraint is OK.
                vv.append(uu[i] + 1)
                move_next = false;
            } else {
                vv.append(0) //Shape constraint is reached.
                move_next = true;//we'll flip the next index.
            }
        } else {
            vv.append(uu[i]) //Nothing to change.
        }
    }
    return ( vv.reduce(true, { $0&&($1 == 0) }) ) ? nil : Array(vv.reversed()); //Returns nil once we reached the Zero Vector.
}

Which gives

add_one_unit(shape:[2,3,3],indexes:[0,0,0]) -> [0,0,1]
add_one_unit(shape:[2,3,3],indexes:[1,2,2]) -> [0,0,0]/nil

Once this is done, this function can be used to enumerate a multilinear map of any shape (a mapping of [i,j,k,...] to a unique index such as matrix to index mapping is necessary and depends on your implementation), or slice a map starting from any particular vector.

Answer 2

Same idea, but less code:

func addOneUnit(shape: [Int], indexes: [Int]) -> [Int]? {
    var next = indexes
    for i in shape.indices.reversed() {
        next[i] += 1
        if next[i] < shape[i] {
            return next 
        }
        next[i] = 0
    }
    return nil
}