Knuth Morris Pratt String matching algorithm

kmp.cpp

#include <iostream>
#include <string>
#include <vector>
using namespace std;

typedef vector<int> vi;
#define pb push_back
#define forn(a, b) for (int a = 0; a < b; ++a)
#define fore(a, b, c) for (int a = b; a < c; ++a)

/**
 * Precomputes the prefix function for the given pattern
 */
vi precompute(string p)
{
    int n = p.length();

    vi T(n, 0);
    T[0] = -1; T[1] = 0; // always

    int sub = 0; // position of substring
    fore(i, 2, n) {
        if (p[i-1] == p[sub]) {
            sub++;
            T[i] = sub;
        } else if (sub > 0) { // there is a mis-match
            sub = T[sub]; // reset sub
        }
    }   

    return T;
}

/**
 * Main Knuth Morris Pratt algorithm
 */
int KMP(string s, string p, vi T) 
{
    int n = s.length(), m = p.length();

    int i = 0, j = 0;

    while (j+i < n) {
        if (s[j+i] == p[i]) {
            if (i == m-1) {    
                return j;
            }
            i++;
        } else {
            j += i - T[i]; // moves the string by i-T[i]
            i = (T[i] > -1) ? T[i] : 0; // moves the pattern to be matched by T[i]
        }
    }

    return n;
}

/**
 * Driver program
 */
int main(int argc, char *argv[]) 
{
    string s, p;
    getline(cin, s);
    getline(cin, p);

    vi T;
    T = precompute(p);

    int k = KMP(s, p, T);

    if (k != s.length()) {
        cout << "Match found at : " << k+1 << endl;       
    } else {
        cout << "No match found!" << endl;
    }
}