Z-Algorithm for pattern matching in strings I was trying to refactor the following Python code (keeping the same time-complexity) which is an implementation of Z-Algorithm for pattern matching in strings.


def compute_z(string):
   n = len(string)
   z = [n]

   L = R = 0
   for i in xrange(1,n):
        if (i &gt; R):
            L = R = i
            while (R &lt; n and string[R-L] == string[R]): R+=1
            z.append(R-L)
            R -= 1
        else:
            k = i-L
            if (z[k] &lt; R-i): z.append(z[k])
            else:
                L = i
                while (R &lt; n and string[R-L] == string[R]): R += 1
                z.append(R-L)
                R -= 1



I was thinking of using lambda but it returns expressions and I need to store the values in z array. So I could not think of a way of using lambda.

The max I could refactor was this:

def g(L,R,string):
    L=i
    while (R &lt; len(string) and string[R-L] == string[R]): R += 1
    return L,R

def compute_z(string):
   n = len(string)
   z = [n]

   L = R = 0
   for i in xrange(1,n):
        if (i &gt; R):
            R = i
            L,R = g(L,R,string)
            z.append(R-L+1)
        else:
            k = i-L
            if (z[k] &lt; R-i): z.append(z[k])  
            else:
                L,R = g(L,R,string)
                z.append(R-L+1)


Could it be any further refactored?

NOTE: Z-Algorithm is defined as z(i) = number of characters in the prefix of s(i) that matches the prefix of S.

For example, a string S = 'abaab'. The substrings of S are

s(0) = 'abaab' = S
s(1) = 'baab'
s(2) = 'aab'
s(3) = 'ab'
s(4) = 'b'


The z values are respectively

z(0) = 5 = length(S)
z(1) = 0
z(2) = 1
z(3) = 2
z(4) = 0


For more details, see this video.
Bugs


In your original compute_z(), you fail to return z.
In your refactored solution, g() has an undefined variable i.
The correct z-values for 'abaab' are [5, 0, 1, 2, 0], not [6, 0, 1, 2, 0].


Questions containing non-working code would normally be closed on Code Review, but since @jonrsharpe has already posted an answer, I'll let it slide.  Since the refactored solution is unreviewable in its current state, I'll examine only the original compute_z().

Reverse engineering

I was, admittedly, unfamiliar with the Z Algorithm.  I stared at compute_z(), befuddled by the cryptically named variables i, k, L, R, and by the purpose of the each if- and else-branch.  Frankly, my experience was closer to reverse-engineering than reading the code.

I came to the realization that I could strip the code down to the following and still get the same result.

def naive_z(string):
    n = len(string)
    z = [n]

    for i in xrange(1, n):
        L = R = i
        while (R &lt; n and string[R - L] == string[R]): R += 1
        z.append(R - L)
    return z


I therefore surmised that the rest of the code served as some kind of optimization, in particular, if (z[k] &lt; R-i): z.append(z[k]).

I also wondered about the R -= 1 statements.  I figured out that there was a loop invariant: at the bottom of the for-loop, string[L : R + 1] was the longest common prefix.

At that point, I decided to give up and start educating myself about the Z Algorithm. ☺

In the end, I decided that the algorithm is complex enough that it wouldn't be reasonable to assume that the underlying principles would be self-evident just by reading the code.  So, I learned something new today.

Observations

Your implementation appears to be a direct translation of the C code at http://codeforces.com/blog/entry/3107.

Superficial criticisms first: you have a few PEP 8 violations.  Notably, the first level of indentation is only 3 spaces, which in Python is a big deal.  The variables L and R violate capitalization conventions, though arguably L is more readable than l.  You also have a few C-isms, such as parentheses around the if-conditions.

Rather than appending one element at a time to z, you should pre-allocate an n-element list at the beginning.  Otherwise, the array might need to be reallocated as it grows.

You miss out on some optimizations due to a botched inequality.  if (z[k] &lt; R-i) should be if z[i] &lt;= R - i.

The R -= 1 statements are awkward, and should be eliminated.  In languages where strings are indexed starting from 0, it would be more idiomatic to use inclusive-exclusive bounds.  In other words, R should point just past the end of the match, such that string[L:R] is the longest matching prefix.

The concern that prompted you to write this question is probably the repetition of the while loops.  It is indeed possible to restructure the code such that the loop only appears once.

def compute_z(string):
    n = len(string)                                                                                                                 
    z = [0] * n                                                                                                                     

    zbox_l, zbox_r, z[0] = 0, 0, n                                                                                                  
    for i in xrange(1, n):          
        if i &lt; zbox_r:              # i is within a zbox                                                                            
            k = i - zbox_l
            if z[k] &lt; zbox_r - i:   
                z[i] = z[k]         # Full optimization                                                                             
                continue            
            zbox_l = i              # Partial optimization                                                                          
        else:
            zbox_l = zbox_r = i     # Match from scratch
        while zbox_r &lt; n and string[zbox_r - zbox_l] == string[zbox_r]:                                                             
            zbox_r += 1 
        z[i] = zbox_r - zbox_l                                                                                                      
    return z