How to split an array into parts where the number of elements will always be different in Swift?

Question

I have an array with shapes of different colors. I need to split an array into multiple arrays without repeating the number of elements. Right now I'm using:

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

The result can sometimes be something like this: [[8],[8],[4]] or [[6],[3],[6]] Need: [[9],[7],[4]] and [[5],[4],[6]]

Edit: This is how the function that returns an array with shapes looks like

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()
    }

Answer 1

This solution isn't exactly short but it seems to work. The tricky part here was generating the sizes of the sub-arrays so I moved that into a separate function.

Basically this function generates random sizes less than what is left of the total size and then add the last bit that is over. It does this in a loop until the number of sub-arrays match size.

private func group(count: Int, into size: Int) -> Set<Int> {
    guard count > size + 2 else { return [] }
    var elementSizes = Set<Int>()

    while true {
        var total = 0
        var limit = count - size

        for _ in 0..<size {
            let value = Int.random(in: 1...limit)
            total += value
            limit -= total
            elementSizes.insert(value)
            if limit <= 0 {
                elementSizes.insert(count - total)
                break
            }
        }
        if elementSizes.count == size { break }
        elementSizes = []
    }
    return elementSizes
}

Once we have this set of sizes the actual splitting could be done like this

func unevenSplit(in size: Int) -> [[Element]] {
    if size < 1 {
        return [self]
    }

    var result = [[Element]]()
    var startIndex = 0
    group(count: count, into: size).forEach {
        result.append(Array(self[startIndex..<(startIndex + $0)]))
        startIndex = $0
    }

    return result
}

Suggestions for improvements are welcome :)

Note my error handling here regarding the value of size is very basic and might need to be improved depending on what kind of values/combinations that can be expected for size and the total array size.

My testing example

let elements = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
let split = elements.unevenSplit(in: 4)

print(split)

Example output:

[[1, 2, 3, 4], [5], [6, 7, 8], [9, 10, 11, 12, 13, 14, 15]]

Answer 2

The usual way to create an random array of unique items is to make a mutable copy of the source array. Then get a random index, add the item at the given index to the result array and remove the item from the temporary array.

For example

let array = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

func randomArray(ofSize size: Int) -> [Int] {
    if size >= array.count { return array.shuffled() }
    var tempArray = array
    return (0..<size).map { _ in tempArray.remove(at: Int.random(in: 0..<tempArray.count)) }
}


randomArray(ofSize: 5)

And you can replace the huge switch statement with

if (2...5).contains(colorsCount) {
    (0..<colorsCount).forEach { i in
        chunkedElements[i] = randomizeFigures(count: chunkedElements[i].count,
                                                     color: colors[i])
    }
    result = Array(chunkedElements.joined())
}