"""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]))