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