EasyLeetCode 28String matching — KMP prefix function

Find the Index of the First Occurrence in a String (LeetCode 28)

Problem 1 of 10 in String Algorithms — Practice Problems

The problem

You get two strings, haystack and needle. Return the index where needle first appears inside haystack, or -1 if it never does.

Input:  haystack = "sadbutsad", needle = "sad"
Output: 0

"sad" appears at index 0 and again at index 6; the first one wins. For haystack = "leetcode", needle = "leeto" the answer is -1.

How to think about it

Start with brute force. Try every start position i and compare needle character by character. That is O(n·m) in the worst case: on haystack = "aaaa…ab" and needle = "aaab", each attempt matches almost the whole needle before failing, then you slide by one and redo all that work.

The key insight: a mismatch tells you what you already matched. If you’ve matched k characters of needle and then fail, those k characters are now known. Some suffix of them may also be a prefix of needle, so a new match may already be under way. Precompute, for every prefix of needle, the length of its longest proper prefix that is also a suffix. That’s the prefix function, stored as lps (“longest prefix that is also a suffix”).

Never move the text pointer back. On a mismatch, don’t restart: set k = lps[k - 1] and try again with the same text character. Each fallback shrinks k, and k only grows by one per text character, so the total work is linear.

Which pattern is this? KMP string matching. The tell is “find a pattern in a text” with a worst case that must be linear. In an interview, haystack.find(needle) is fine to mention, but KMP is what they’re testing.

Approach

  1. If needle is empty, return 0.
  2. Build lps over needle: keep k as the current matched-prefix length. For each i, fall back with k = lps[k - 1] while needle[i] != needle[k]; if they match, k += 1; store lps[i] = k.
  3. Scan haystack with the same loop, but compare haystack[i] with needle[k].
  4. When k reaches len(needle), the match ends at i, so return i - k + 1.
  5. If the scan finishes, return -1.

Solution

def strStr(haystack, needle):
    if not needle:
        return 0
    lps = [0] * len(needle)
    k = 0
    for i in range(1, len(needle)):              # build the prefix function
        while k and needle[i] != needle[k]:
            k = lps[k - 1]                       # fall back to a shorter border
        if needle[i] == needle[k]:
            k += 1
        lps[i] = k
    k = 0
    for i in range(len(haystack)):               # scan: i never moves backwards
        while k and haystack[i] != needle[k]:
            k = lps[k - 1]
        if haystack[i] == needle[k]:
            k += 1
        if k == len(needle):
            return i - k + 1
    return -1
# strStr("sadbutsad", "sad") -> 0 ;  strStr("leetcode", "leeto") -> -1

Dry run

The simple example matches at once, so here is one where the fallback matters: haystack = "aabaaabaaac", needle = "aabaaac". The code builds lps = [0, 1, 0, 1, 2, 2, 0].

i haystack[i] k before k after fallback k after
0–5 a a b a a a 0 → 5 no fallback 6
6 b 6 2 (lps[5]) 3
7 a 3 3 4
8 a 4 4 5
9 a 5 5 6
10 c 6 6 7, full match

At i = 6 the needle wanted c but saw b. Instead of restarting, k drops to 2 ("aa" is still matched) and continues. The match ends at 10, so the answer is 10 - 7 + 1 = 4.

Complexity

  • Time: O(n + m). Building lps is O(m) and the scan is O(n): k rises at most once per character, and each fallback lowers it.
  • Space: O(m). For the lps table.

Edge cases and pitfalls

  • Empty needle returns 0 by convention.
  • needle longer than haystack returns -1; the scan never reaches full length.
  • Using while vs if for the fallback. It must be a while: one fallback may not be enough.
  • Off-by-one in the answer. When k == len(needle) at index i, the start is i - k + 1, not i - k.
  • Starting the lps loop at i = 0 would match needle[0] with itself and give lps[0] = 1. Start at 1.