Calculate shape of N-dimensional array from a flat array

Using Swift
,
1

I have a flat array as shown below that represents a 3D array. The example prints the shape of the array as (2, 3, 4) which means there are two arrays each with a size of 3-by-4 elements.

let arr = [[[1, 2, 3, 4],
            [5, 6, 7, 8],
            [9, 10, 11, 12]],

           [[13, 14, 15, 16],
            [17, 18, 19, 20],
            [21, 22, 23, 24]]]

let shape = (arr.count, arr[0].count, arr[0][0].count)
print("shape is", shape)

I would like to calculate the shape of this array (or any N-dimensional array) using a function. My failed attempt is shown below.

func getShape<T>(_ arrays: [T]...) -> [Int] {
    if type(of: arrays) != Array<T>.self {
        return [arrays.count]
    }

    var shape = [Int]()
    for subArray in arrays {
        shape.append(subArray.count)
    }
    return shape
}

let s = getShape(arr)
print("shape is", s)

The first problem seems to be that the variadic input parameter [T]... is not being iterated over. It is interpreted as a single array and not an array of arrays. I also think the function needs to be recursive to properly calculate the shape. But before I can do that, I need to figure out how to properly define the input parameter. Does anyone have suggestions on how to calculate the shape of a flat array that represents an N-dimensional array?

2

You cannot check if the type of an array element is itself an array, because in doing so you need to specify the inner array's element. You can, however, use a protocol that gives you count (such as Collection or BidirectionalCollection), and recursively define your shape function as follows:

func shape(_ arr: some Collection) -> [Int] {
  if let first = arr.first as? any Collection {
    return [arr.count] + shape(first)
  } else {
    return [arr.count]
  }
}
5 Likes
3

Ah, good idea. I noticed that using any such as func shape(_ arr: any Collection) also works. So what's the difference between using any or some for the function input?

4

any Collection will box up the argument in an existential to guarantee the size of the parameter since the compiler doesn't know how much space it needs. Technically, you don't need that in this case so some Collection will definitely give the compiler enough hints to be more performant.

And yet, you immediately try to pull the first parameter of the collection out as an any Collection so you already have an existential for the recursive call. I'm not sure if it you really gain anything using some Collection then.

5

So I can do the following which correctly prints a shape of 2x3.

func getShape(_ arr: some Collection) -> [Int] {
    if let first = arr.first as? any Collection {
        return [arr.count] + getShape(first)
    } else {
        return [arr.count]
    }
}

let s = getShape([[1, 2, 3], [4, 5, 6]])
print("shape is", s)

shape is [2, 3]

I would like to use this with a shaped array structure that stores its data as a flat array:

struct ShapedArray<T> {
    let shape: [Int]
    let data: [T]

    init(shape: [Int], array: [T]) {
        self.shape = shape
        self.data = array
    }

    init(arrays: [T]...) {
        self.shape = getShape(arrays)
        self.data = arrays.flatMap { $0 }
    }
}

The first init method takes the shape and a flat array. The second init method determines the shape from some N-dimensional array. But the second init method does not flatten the array for the underlying data storage.

let sa = ShapedArray(arrays: [[1, 2, 3], [4, 5, 6]])
print("shape is", sa.shape)
print("data is", sa.data)

shape is [1, 2, 3]
data is [[1, 2, 3], [4, 5, 6]]

The results I expect to see are:

shape is [2, 3]
data is [1, 2, 3, 4, 5, 6]

I tried to use a collection instead of a variadic parameter but apparently a collection can't be flattened to an array:

init(arrays: some Collection) {
    self.shape = getShape(arrays)
    self.data = arrays.flatMap { $0 }  // This doesn't work
}

Is there a way I can use the getShape() function within the second init method and flatten the input?

6

The (arrays: [T]...) signature means that you’re expected to pass multiple [T] parameters, like ShapedArray(arrays: [1], [1, 2, 3]). What you need to do here to recursively flatten the array is to update getShape to return an array of Int representing the shape of the array, and a flattened array of values. But I’m not sure if there’s a way to specify that in the type system since there isn’t really a way to group [T] and [[T]] and [[[T]]] as the same type where you can do something with T.

2 Likes
7

Ok, I thought [T]... meant multiple subarrays such as [[T]] or [[[T]]] but if it behaves as multiple [T] such as [T], [T], [T] then that would explain the results I got above. If I use some Collection instead of [T]..., is there a way to flatten the collection to an array?

8

Yes, but you would not be able to constrain the deeply-nested-element type to T:

func flatten(_ arr: [some Any]) -> [any Any] {
  if let arr = arr as? [[any Any]] {
    return flatten(arr.reduce([], +))
  } else {
    return arr
  }
}
9

Well that's not encouraging. I'll post another topic related to flattening with type information. But for now I'm able to get the shape of an N-dimensional array which is what this topic is originally about. Thanks for the help.

10

Imagine having this:

:star_struck:

11

That would be a nice addition to Swift. Do you think the proposal will get accepted and how long would it take to get released as a Swift feature?

12

This is a future direction of a pitch, so ... :D

I just wanted to let you know that this idea appeared somewhere, so you could follow the conversation if you'd like :)