在 Swift 中从嵌套数组创建 ShapedArray

Question

我正在尝试创建一个形状数组结构来处理 Swift 中的 N 维（多维）数值数据。我想使用嵌套数组（数组的数组）创建结构，如下所示：

// This would create a double-precision 2D array with a shape of 2x3
// ⎛ 1, 2, 3 ⎞
// ⎝ 4, 5, 6 ⎠
let arr = ShapedArray<Double>([[1, 2, 3], [4, 5, 6]])

// This would create a single-precision 3D array with a shape of 2x2x3
// ⎛ ⎛ 1, 2, 3 ⎞ ⎞
// ⎜ ⎝ 4, 5, 6 ⎠ ⎟
// ⎜ ⎛ 7, 8, 9 ⎞ ⎟
// ⎝ ⎝ 0, 1, 2 ⎠ ⎠
let arr = ShapedArray<Float>([[[1, 2, 3],
                               [4, 5, 6]],
                              [[7, 8, 9],
                               [0, 1, 2]]])

以下是我实现这一点的尝试。

ShapedArray

结构将数据存储在平面数组中。

getShape

函数返回嵌套数组的形状。

flatten

函数从嵌套数组返回单个数组（展平数组）。

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

func flatten(_ arrays: [Any]) -> [Any] {
    var result = [Any]()

    for val in arrays {
        if let arr = val as? [Any] {
            result.append(contentsOf: flatten(arr))
        } else {
            result.append(val)
        }
    }

    return result
}

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

    init(arrays: [Any]) {
        self.shape = getShape(arrays)
        self.data = flatten(arrays) as! [T]
    }
}

使用

ShapedArray

结构的一些示例如下所示。

// This works for 1D array
let sa = ShapedArray<Int>(arrays: [1, 2, 3, 4])

// This doesn't work
// let sa = ShapedArray<Float>(arrays: [1, 2, 3, 4])

// This works for 2D array
let saa = ShapedArray<Int>(arrays: [[0, 1, 2], [3, 4, 5]])

// This doesn't work
// let saa = ShapedArray<Float>(arrays: [[0, 1, 2], [3, 4, 5]])

// This works for 3D array
let saaa = ShapedArray<Int>(arrays: [[[0, 1, 2],
                                      [3, 4, 5]],
                                     [[6, 7, 8],
                                      [9, 0, 1]]])

此代码仅适用于整数数组。如果我尝试创建一个

ShapedArray

的浮点数或双精度数，我会收到如下错误：

Could not cast value of type 'Swift.Int' (0x7ff84ad1b2a0) to 'Swift.Float' (0x7ff84ad1ae28).

因此显然

ShapedArray

的类型信息没有传递给展平数组的结果。如何展平嵌套数组并定义元素类型？或者我应该使用不同的方法？

Answer 1

问题是你正在使用

[Any]

，所以 Swift 不知道数组文字代表

Float

/

Double

的嵌套数组。由于您使用整数文字作为数组元素，这就是 Swift 认为您想要创建的内容。然后，您将这些

Int

转换为

Float

/

Double

，从而导致崩溃。

这可以简单地通过注释数组文字的类型来修复：

let shapedArray = ShapedArray<Float>(arrays: 
     [[[0, 1, 2],
       [3, 4, 5]],
      [[6, 7, 8],              
       [9, 0, 1]]] as [[[Float]]]
)

也就是说，使用

[Any]

并不是很安全。一个更类型安全的解决方案是使用 sum 类型的数组，如下所示：

indirect enum ArrayOrElement<T>: ExpressibleByArrayLiteral {
    case array([ArrayOrElement<T>])
    case element(T)
    
    init(arrayLiteral elements: T...) {
        self = .array(elements.map { .element($0) })
    }
}

// ExpressibleByXXXLiteral conformances
extension ArrayOrElement: ExpressibleByIntegerLiteral where T: ExpressibleByIntegerLiteral {
    init(integerLiteral value: T.IntegerLiteralType) {
        self = .element(T(integerLiteral: value))
    }
}

extension ArrayOrElement: ExpressibleByFloatLiteral where T: ExpressibleByFloatLiteral {
    init(floatLiteral value: T.FloatLiteralType) {
        self = .element(T(floatLiteral: value))
    }
}
// you can also do this for boolean and string literals...

然后您可以编写适用于此类型的

getShape

和

flatten

：

func getShape<T>(_ arr: [ArrayOrElement<T>]) -> [Int] {
    switch arr.first {
    case .array(let inner):
        [arr.count] + getShape(inner)
    default:
        [arr.count]
    }
}

func flatten<T>(_ arrays: [ArrayOrElement<T>]) -> [T] {
    var result = [T]()
    for val in arrays {
        switch val {
        case .array(let arr):
            result.append(contentsOf: flatten(arr))
        case .element(let e):
            result.append(e)
        }
    }

    return result
}

然后

ShapedArray.init

可以拍

[ArrayOrElement<T>]

。

struct ShapedArray<T>: ExpressibleByArrayLiteral {
    let shape: [Int]
    let data: [T]
    
    init(arrayLiteral elements: ArrayOrElement<T>...) {
        self.shape = getShape(elements)
        self.data = flatten(elements)
    }
}

用途：

let shapedArray: ShapedArray<Double> = [
    .array([
        [0, 1, 2],
        [3, 4, 5],
    ]),
    .array([
        [6, 7, 8],
        [9, 0, 1],
    ]),
]

/*
The ExpressibleByXXXLiteral conformance might suggest that you could also remove the ".array(...)" calls like this:

let shapedArray: ShapedArray<Double> = [
    [
        [0, 1, 2],
        [3, 4, 5],
    ],
    [
        [6, 7, 8],
        [9, 0, 1],
    ],
]

but this doesn't work. Swift's type inference algorithm seems to have
reached its limits when you have triple nested arrays.
*/