Edit Distance in Swift

EditDistance.swift

enum EditOperation<Element> {
    case Delete(index: Int)
    case Insert(index: Int, element: Element)
    case Substitute(index: Int, element: Element)
    case DoNothing

    var distance: Int {
        switch self {
        case .DoNothing:
            return 0
        default:
            return 1
        }
    }
}

class EditOperationChain<Element> {
    typealias ConcreteEditOperation = EditOperation<Element>
    typealias ConcreteEditOperationChain = EditOperationChain<Element>

    private(set) var current: ConcreteEditOperation
    private(set) var previousChain: ConcreteEditOperationChain? = nil
    private(set) var distance: Int = 0

    init(current: ConcreteEditOperation, previousChain: ConcreteEditOperationChain? = nil) {
        self.current = current
        self.previousChain = previousChain

        self.distance = (self.previousChain?.distance ?? 0) + current.distance
    }
}


class Matrix<Value> {

    private(set) var columns: Int
    private(set) var rows: Int

    private var backstage: [Value?]

    init(columns: Int, rows: Int) {
        self.columns = columns
        self.rows = rows
        self.backstage = Array(count: columns * rows, repeatedValue: nil)
    }

    subscript(column: Int, row: Int) -> Value? {
        get {
            return self.backstage[self.columns * row + column]
        }
        set {
            self.backstage[self.columns * row + column] = newValue
        }
    }
}


class MinumumEditDistanceCalculator<Element: Equatable> {
    typealias ConcreteEditOperation = EditOperation<Element>
    typealias ConcreteEditOperationChain = EditOperationChain<Element>

    var before: [Element]
    var after: [Element]
    var matrix: Matrix<EditOperationChain<Element>>!

    init(before: [Element], after: [Element]) {
        self.before = before
        self.after = after
    }

    func solve() -> [EditOperation<Element>] {
        var matrix = Matrix<ConcreteEditOperationChain>(columns: self.before.count + 1, rows: self.after.count + 1)

        matrix[0, 0] = ConcreteEditOperationChain(
            current: .DoNothing
        )

        // Deletions
        if self.before.count > 0 {
            for i in 1...self.before.count {
                matrix[i, 0] = ConcreteEditOperationChain(
                    current: .Delete(index: 0/*i - 1*/),
                    previousChain: matrix[i - 1, 0]
                )
            }
        }

        // Insertions
        if self.after.count > 0 {
            for j in 1...self.after.count {
                matrix[0, j] = ConcreteEditOperationChain(
                    current: .Insert(index: 0/*j - 1*/, element: self.after[j - 1]),
                    previousChain: matrix[0, j - 1]
                )
            }
        }

        if self.before.count > 0 && self.after.count > 0 {
            for i in 1...self.before.count {
                for j in 1...self.after.count {

                    if self.before[i - 1] == self.after[j - 1] {
                        matrix[i, j] = matrix[i - 1, j - 1]
                        continue
                    }

                    let variants = [
                        matrix[i - 1, j]!,  // deletion
                        matrix[i, j - 1]!,  // insertion
                        matrix[i - 1, j - 1]!  // substitution
                    ] as [ConcreteEditOperationChain]

                    let distances = variants.map { $0.distance }
                    let minimumIndex = distances.indexOf(distances.minElement()!)!
                    let minimumEditChain = variants[minimumIndex]

                    var editChain: ConcreteEditOperationChain? = nil
                    switch minimumIndex {
                    case 0:
                        // Deletion
                        editChain = ConcreteEditOperationChain(current: .Delete(index: j), previousChain: minimumEditChain)
                        break
                    case 1:
                        // Insertion
                        editChain = ConcreteEditOperationChain(current: .Insert(index: j - 1, element: self.after[j - 1]), previousChain: minimumEditChain)
                        break
                    case 2:
                        // Substitution
                        editChain = ConcreteEditOperationChain(current: .Substitute(index: j - 1, element: self.after[j - 1]), previousChain: minimumEditChain)
                        break
                    default:
                        assertionFailure()
                    }

                    matrix[i, j] = editChain
                }
            }
        }

        var currentChain: ConcreteEditOperationChain? = matrix[self.before.count, self.after.count]

        var editOperations: [ConcreteEditOperation] = []

        while currentChain != nil {
            let editOperation = currentChain!.current

            switch editOperation {
            case .DoNothing:
                break
            default:
                editOperations.insert(editOperation, atIndex: 0)
                break
            }

            currentChain = currentChain!.previousChain
        }

        return editOperations
    }
}