Leetcode (Python): Implement strStr()

Implement strStr().

Returns a pointer to the first occurrence of needle in haystack, or null if needle is not part of haystack.

Solution:

We first preprocess the needle, in such a way that we obtain the equivalent states, for instance:

   a a b a a b c
   0 0 1 0 1 2 3

That means that if we are have found aabaab, we have also found aab. When we process the haystack, when we arrive to a state that we cannot go on, we change to its previous equivalent state.

class Solution:
    # @param haystack, a string
    # @param needle, a string
    # @return a string or None
    def strStr(self, haystack, needle):
        if len(needle)==0:
            return haystack
        equivalentState=[0]
        previousState = 0;
        for i in range(1,len(needle)):
            equivalentState.append(previousState)
            if needle[previousState] == needle [i]:
                previousState += 1
            else:
                previousState = 0
        position = 0
        state = 0
        while position < len(haystack):
            if haystack[position] == needle[state]:
                position += 1
                state += 1
                if state == len(needle):
                    return haystack[position-state:]
            else:
                if state == 0:
                    position +=1
                else:
                    state = equivalentState[state]
        return None