Swift实现数组拆分：各子数组元素数量互不相同的方法

我有一个包含不同颜色形状元素的数组，需要将它拆分成多个子数组，要求每个子数组的元素数量互不重复。目前我用下面的代码实现拆分：

extension Array {
    func split(into size: Int) -> [[Element]] {
            return (0..<size).map {
                stride(from: $0, to: count, by: size).map { self[$0] }
            }
        }
}

但拆分结果偶尔会出现子数组元素数量重复的情况，比如[[8],[8],[4]]或[[6],[3],[6]]，而我想要的是所有子数组元素数量都不重复的结果，例如[[9],[7],[4]]和[[5],[4],[6]]。

fileprivate func getTreasureElements(totalCount: Int,
                                     colors: [MysteriosTreasureElementColor],
                                     colorsCount: Int,
                                     figuresCount: Int) -> [MysteriosTreasureElement] {

    var elements: [MysteriosTreasureElement] = []
    let circleFigure = MysteriosTreasureElementFigure.circle
    let randomColor = MysteriosTreasureElementColor.allCases.randomElement()!
    let randomFigure = MysteriosTreasureElementFigure.allCases.randomElement()!

    let newElement = MysteriosTreasureElement(icon: UIImage(named: "treasureIcon_\(randomColor)_\(circleFigure)")!,
                                           color: randomColor,
                                           figure: randomFigure)
    for _ in 0...totalCount+1 {
        elements.append(newElement)
    }
    var chunkedElements = elements.split(into: colorsCount)
    var result: [MysteriosTreasureElement] = []
    switch colorsCount {
    case 2:
        
        chunkedElements[0] = randomizeFigures(count: chunkedElements[0].count,
                                              color: colors[0])
        
        chunkedElements[1] = randomizeFigures(count: chunkedElements[1].count,
                                              color: colors[1])
        result = Array(chunkedElements.joined())
    case 3:
        chunkedElements[0] = randomizeFigures(count: chunkedElements[0].count,
                                              color: colors[0])
        chunkedElements[1] = randomizeFigures(count: chunkedElements[1].count,
                                              color: colors[1])
        chunkedElements[2] = randomizeFigures(count: chunkedElements[2].count,
                                              color: colors[2])
        result = Array(chunkedElements.joined())
        
    case 4:
        chunkedElements[0] = randomizeFigures(count: chunkedElements[0].count,
                                              color: colors[0])
        chunkedElements[1] = randomizeFigures(count: chunkedElements[1].count,
                                              color: colors[1])
        chunkedElements[2] = randomizeFigures(count: chunkedElements[2].count,
                                              color: colors[2])
        chunkedElements[3] = randomizeFigures(count: chunkedElements[3].count,
                                              color: colors[3])
        result = Array(chunkedElements.joined())
        
    case 5:
        chunkedElements[0] = randomizeFigures(count: chunkedElements[0].count,
                                              color: colors[0])
        chunkedElements[1] = randomizeFigures(count: chunkedElements[1].count,
                                              color: colors[1])
        chunkedElements[2] = randomizeFigures(count: chunkedElements[2].count,
                                              color: colors[2])
        chunkedElements[3] = randomizeFigures(count: chunkedElements[3].count,
                                              color: colors[3])
        chunkedElements[4] = randomizeFigures(count: chunkedElements[4].count,
                                              color: colors[4])
        
        result = Array(chunkedElements.joined())
        
    default:
        break
    }
    while result.count > totalCount {
        result.removeLast()
    }
    return result.shuffled()
}

内容的提问来源于stack exchange，提问作者Avoopihra