Simple implementation of Knuth-Morris-Pratt (KMP) algorithm in Swift

Main.swift

//                .─────────────────────.
//          _.───'                       `────.
//        ,'                                   `.
//       (     init(next/1: ℕ[m]; P/1: ∑[m])     )
//        '─.                                 ,─'
//           `────.                     _.───'
//                 `──────────┬────────'
//                            │
//  ┌─────────────────────────┴──────────────────────┐
//  │           next[1] := i := 0; j := 1            │
//  ├────────────────────────────────────────────────┤
//  │                      j < m                     │
//  │   ┌────────────────────────────────────────────┤
//  │   │\             P[i+1] = P[j+1]              /│
//  │   ├──────────────┬─────────────────────────────┤
//  │   │     i++      │\           i = 0           /│
//  │   ├──────────────┼──────────────┬──────────────┤
//  │   │     j++      │     j++      │              │
//  │   ├──────────────┼──────────────┤ i := next[i] │
//  │   │ next[j] := i │ next[j] := 0 │              │
//  └───┴──────────────┴──────────────┴──────────────┘
func initNext(next: inout [Int], pattern: String) {
    var i = 0
    var j = 0
    next = [Int](repeating: i, count: m)
    
    while j + 1 < m {
        if (pattern[i] == pattern[j+1]) {
            i += 1
            j += 1
            next[j] = i
        } else {
            if i == 0 {
                j += 1
                next[j] = 0
            } else {
                i = next[i - 1]
            }
        }
    }
}

//               .─────────────────────.
//         _.───'                       `────.
//       ,'                                   `.
//      (   KMP(T/1: ∑[n]; P/1: ∑[m]; S: ℕ{})   )
//       '─.                                 ,─'
//          `────.                     _.───'
//                `─────────┬─────────'
//                          │
// ┌────────────────────────┴─────────────────────────┐
// │           next/1: ℕ[m]; init(next, P)            │
// ├──────────────────────────────────────────────────┤
// │               S := {}; i := j := 0               │
// ├──────────────────────────────────────────────────┤
// │                      i < n                       │
// │   ┌──────────────────────────────────────────────┤
// │   │\              T[i+1] = P[j+1]               /│
// │   ├─────────────────────────┬────────────────────┤
// │   │        i++; j++         │\      j = 0       /│
// │   ├─────────────────────────┼─────┬──────────────┤
// │   │\         j = m         /│     │              │
// │   ├──────────────────┬──────┤     │              │
// │   │ S := S U {i - m} │      │ i++ │ j := next[j] │
// │   ├──────────────────┤ SKIP │     │              │
// │   │   J := next[m]   │      │     │              │
// └───┴──────────────────┴──────┴─────┴──────────────┘
func kmp(text: String, pattern: String, solutions: inout [Int]) {
    var next = [Int]()
    initNext(next: &next, pattern: pattern)
    solutions = [Int]()
    var i = 0
    var j = 0
    while i < n {
        if text[i] == pattern[j] {
            i += 1
            j += 1
            if j == m {
                solutions.append(i - m)
                j = next[m-1]
            }
        } else {
            if j == 0 {
                i += 1
            } else {
                j = next[j-1]
            }
        }
    }
}

// -------------------------------------
// --------------- Usage ---------------
// -------------------------------------

let text = "ABABABBABABBABABA"
let pattern = "ABABBABA"
let n = text.count
let m = pattern.count

var solutions = [Int]()
kmp(text: text, pattern: pattern, solutions: &solutions)
print("Solutions: \(solutions)")

SharedCode.swift

public extension String {
    
    var length: Int {
        return count
    }
    
    subscript (i: Int) -> String {
        return self[i ..< i + 1]
    }
    
    func substring(fromIndex: Int) -> String {
        return self[min(fromIndex, length) ..< length]
    }
    
    func substring(toIndex: Int) -> String {
        return self[0 ..< max(0, toIndex)]
    }
    
    subscript (r: Range<Int>) -> String {
        let range = Range(uncheckedBounds: (lower: max(0, min(length, r.lowerBound)),
                                            upper: min(length, max(0, r.upperBound))))
        let start = index(startIndex, offsetBy: range.lowerBound)
        let end = index(start, offsetBy: range.upperBound - range.lowerBound)
        return String(self[start ..< end])
    }
    
}