I recently completed 10 Days of Statistics on HackerRank. Yay!
In the process, I needed some matrix operations for a medium-difficulty problem. And here they are, code style be damned :
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
func transpose(_ matrix: [[Double]]) -> [[Double]] { let rowCount = matrix.count let colCount = matrix[0].count var transposed : [[Double]] = Array(repeating: Array(repeating: 0.0, count: rowCount), count: colCount) for rowPos in 0..<matrix.count { for colPos in 0..<matrix[0].count { transposed[colPos][rowPos] = matrix[rowPos][colPos] } } return transposed } func multiply(_ A: [[Double]], _ B: [[Double]]) -> [[Double]] { let rowCount = A.count let colCount = B[0].count var product : [[Double]] = Array(repeating: Array(repeating: 0.0, count: colCount), count: rowCount) for rowPos in 0..<rowCount { for colPos in 0..<colCount { for i in 0..<B.count { product[rowPos][colPos] += A[rowPos][i] * B[i][colPos] } } } return product } // gauss jordan inversion func inverse(_ matrix: [[Double]]) -> [[Double]] { // augment matrix var matrix = matrix var idrow = Array(repeating: 0.0, count: matrix.count) idrow[0] = 1.0 for row in 0..<matrix.count { matrix[row] += idrow idrow.insert(0.0, at:0) idrow.removeLast() } // partial pivot for row1 in 0..<matrix.count { for row2 in row1..<matrix.count { if abs(matrix[row1][row1]) < abs(matrix[row2][row2]) { (matrix[row1],matrix[row2]) = (matrix[row2],matrix[row1]) } } } // forward elimination for pivot in 0..<matrix.count { // multiply let arg = 1.0 / matrix[pivot][pivot] for col in pivot..<matrix[pivot].count { matrix[pivot][col] *= arg } // multiply-add for row in (pivot+1)..<matrix.count { let arg = matrix[row][pivot] / matrix[pivot][pivot] for col in pivot..<matrix[row].count { matrix[row][col] -= arg * matrix[pivot][col] } } } // backward elimination for pivot in (0..<matrix.count).reversed() { // multiply-add for row in 0..<pivot { let arg = matrix[row][pivot] / matrix[pivot][pivot] for col in pivot..<matrix[row].count { matrix[row][col] -= arg * matrix[pivot][col] } } } // remove identity for row in 0..<matrix.count { for _ in 0..<matrix.count { matrix[row].remove(at:0) } } return matrix } let X = [ [1.0, 2.0, 3.0], [4.0, 5.0, 11.0], [7.0, 8.0, 9.0] ] let XI = inverse(X) let I = multiply(X,XI) print(I) |
That’s the identity matrix popping out at the end, which validates my implementation.
But what’s this? A 60-line method? Uncle Bob would not be pleased.
Comments should not take the place of good variable/method names. Those section comments give clues as to where my methods should be :
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
func augment(_ matrix: [[Double]]) -> [[Double]] { var augmented = matrix var idrow = Array(repeating: 0.0, count: matrix.count) idrow[0] = 1.0 for row in 0..<matrix.count { augmented[row] += idrow idrow.insert(0.0, at:0) idrow.removeLast() } return augmented } func deaugment(_ matrix: [[Double]]) -> [[Double]] { var deaugmented = matrix for row in 0..<matrix.count { for _ in 0..<matrix.count { deaugmented[row].remove(at:0) } } return deaugmented } func partialPivot(_ matrix: inout [[Double]]) { for row1 in 0..<matrix.count { for row2 in row1..<matrix.count { if abs(matrix[row1][row1]) < abs(matrix[row2][row2]) { (matrix[row1],matrix[row2]) = (matrix[row2],matrix[row1]) } } } } func scaleRow(_ matrix: inout [[Double]], row: Int, scale: Double) { for col in 0..<matrix[row].count { matrix[row][col] *= scale } } func addRow(_ matrix: inout [[Double]], row: Int, scaledBy: Double, toRow: Int) { for col in 0..<matrix[row].count { matrix[toRow][col] += scaledBy * matrix[row][col] } } func pivot(_ matrix: inout [[Double]], row pivotRow: Int, col pivotCol: Int, forward: Bool) { let scale = 1.0 / matrix[pivotRow][pivotCol] scaleRow(&matrix, row: pivotRow, scale: scale) if forward { for toRow in (pivotRow+1)..<matrix.count { let scaleBy = -1.0 * matrix[toRow][pivotCol] addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow) } } else { for toRow in (0..<pivotRow).reversed() { let scaleBy = -1.0 * matrix[toRow][pivotCol] addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow) } } } func gaussJordanInverse(_ matrix: [[Double]]) -> [[Double]] { var matrix = augment(matrix) partialPivot(&matrix) for p in 0..<matrix.count { pivot(&matrix, row: p, col: p, forward: true) } for p in (0..<matrix.count).reversed() { pivot(&matrix, row: p, col: p, forward: false) } matrix = deaugment(matrix) return matrix } let X = [ [1.0, 2.0, 3.0], [4.0, 5.0, 11.0], [7.0, 8.0, 9.0] ] let XI = gaussJordanInverse(X) let I = multiply(X,XI) print(I) |
Better. Uncle Bob would be proud (or give me credit for trying, anyway).
One more thing : Thanks to StackOverflow user Alexander, I have an even better way to express that pivot loop :
|
1 2 3 4 5 6 7 8 9 10 11 |
func pivot(_ matrix: inout [[Double]], row pivotRow: Int, col pivotCol: Int, forward: Bool) { let scale = 1.0 / matrix[pivotRow][pivotCol] scaleRow(&matrix, row: pivotRow, scale: scale) let range = forward ? AnyCollection((pivotRow+1)..<matrix.count) : AnyCollection((0..<pivotRow).reversed()) for toRow in range { let scaleBy = -1.0 * matrix[toRow][pivotCol] addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow) } } |
