Skip to content

Instantly share code, notes, and snippets.

@suyash
Created July 8, 2020 13:26
Show Gist options
  • Save suyash/09f8b2e49abd415f15f3838420ce7534 to your computer and use it in GitHub Desktop.
Save suyash/09f8b2e49abd415f15f3838420ce7534 to your computer and use it in GitHub Desktop.
Knuth-Morris-Pratt implementation
//! Knuth-Morris-Pratt implementation
//!
//! See Chapter 17 in Blandy + Orendorff
/// computes the KMP longest prefix-suffix function
pub fn kmp<T: AsRef<[u8]>>(pattern: T) -> Vec<usize> {
let pattern = pattern.as_ref();
let n = pattern.len();
let mut ans = vec![0; n + 1];
for (i, c) in pattern.iter().enumerate().skip(1) {
let mut j = ans[i];
loop {
if c == &pattern[j] {
ans[i + 1] = j + 1;
break;
}
if j == 0 {
break;
}
j = ans[j];
}
}
ans
}
/// find pattern in haystack
pub fn find<T: AsRef<[u8]>>(pattern: T, haystack: T) -> Vec<usize> {
let pattern = pattern.as_ref();
let haystack = haystack.as_ref();
let mut ans = Vec::new();
let f = kmp(pattern);
let mut i = 0;
let mut j = 0;
while i < haystack.len() {
if haystack[i] == pattern[j] {
i += 1;
j += 1;
if j == pattern.len() {
ans.push(i - pattern.len());
j = f[j];
}
} else if j > 0 {
j = f[j];
} else {
i += 1;
}
}
ans
}
#[cfg(test)]
mod tests {
#[test]
fn test_kmp_fn() {
let p = super::kmp("ABABAC");
assert_eq!(p, vec![0, 0, 0, 1, 2, 3, 0]);
}
#[test]
fn test_find() {
assert_eq!(
super::find("ABABA", "ABABABABABABABABABABA"),
vec![0, 2, 4, 6, 8, 10, 12, 14, 16]
);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment