64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
"""Z-algorithm: linear-time Z-array and substring search.
|
|
|
|
Z[i] = length of the longest substring of s starting at i that is also a
|
|
prefix of s. Z[0] is conventionally len(s) (or 0, depending on source);
|
|
here Z[0] = len(s) for convenience, matching the common competitive-programming
|
|
convention.
|
|
|
|
Complexity: O(n) time, O(n) space.
|
|
"""
|
|
|
|
|
|
def z_function(s: str):
|
|
"""Return the Z-array of s."""
|
|
n = len(s)
|
|
z = [0] * n
|
|
if n == 0:
|
|
return z
|
|
z[0] = n
|
|
l = r = 0
|
|
for i in range(1, n):
|
|
if i < r:
|
|
z[i] = min(r - i, z[i - l])
|
|
while i + z[i] < n and s[z[i]] == s[i + z[i]]:
|
|
z[i] += 1
|
|
if i + z[i] > r:
|
|
l, r = i, i + z[i]
|
|
return z
|
|
|
|
|
|
def search(pattern: str, text: str):
|
|
"""Return list of starting indices in text where pattern occurs (possibly overlapping)."""
|
|
if pattern == "":
|
|
return list(range(len(text) + 1))
|
|
if len(pattern) > len(text):
|
|
return []
|
|
combined = pattern + "\x00" + text
|
|
z = z_function(combined)
|
|
m = len(pattern)
|
|
res = []
|
|
offset = m + 1
|
|
for i in range(offset, len(combined)):
|
|
if z[i] >= m:
|
|
res.append(i - offset)
|
|
return res
|
|
|
|
|
|
def _self_test():
|
|
assert z_function("") == []
|
|
assert z_function("a") == [1]
|
|
assert z_function("aaaa") == [4, 3, 2, 1]
|
|
assert z_function("ababa") == [5, 0, 3, 0, 1]
|
|
assert search("ab", "abxab") == [0, 3]
|
|
assert search("aa", "aaaa") == [0, 1, 2]
|
|
assert search("", "abc") == [0, 1, 2, 3]
|
|
assert search("xyz", "abc") == []
|
|
print("OK z_algorithm self-test")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_self_test()
|
|
import sys
|
|
if len(sys.argv) > 1:
|
|
print(z_function(sys.argv[1]))
|