프로그래머스 - 정사각형으로 만들기

func solution(_ arr:[[Int]]) -> [[Int]] {
    var newArr = arr
    let maximum = max(arr.map { $0.count }.max()!, arr.count)
    for i in 0..<maximum { // 0,1,2,3
        if newArr.count > i {
            while newArr[i].count != maximum { // 행이 열보다 많을 때
                newArr[i].append(0)
            }
        } else {
            var tmpArr: [Int] = Array(repeating: 0, count: maximum)
            newArr.append(tmpArr)
        }
    }
    
    return newArr
}