From 916aaee6307998896835da0b5cc65c2b8f88c723 Mon Sep 17 00:00:00 2001
From: Tommy Li <dev@stdx.space>
Date: Sat, 31 Aug 2024 19:03:25 +0800
Subject: [PATCH] feat: add z function

---
 string.hpp      | 20 ++++++++++++++++++++
 string_test.cpp |  8 ++++++++
 2 files changed, 28 insertions(+)

diff --git a/string.hpp b/string.hpp
index 87d7d08..2bb527e 100644
--- a/string.hpp
+++ b/string.hpp
@@ -32,6 +32,26 @@ std::vector<int> prefix_function(R &&s) {
 	return pi;
 }
 
+template <std::ranges::random_access_range R>
+    requires std::ranges::sized_range<R>
+std::vector<int> z_function(R &&s) {
+	std::vector<int> z(std::ranges::size(s));
+	int l = 0, r = 0;
+	for (int i = 1; i < std::ranges::ssize(s); i++) {
+		if (i < r) {
+			z[i] = std::min(r - i, z[i - l]);
+		}
+		while (i + z[i] < std::ranges::ssize(s) && s[z[i]] == s[i + z[i]]) {
+			z[i]++;
+		}
+		if (i + z[i] > r) {
+			l = i;
+			r = i + z[i];
+		}
+	}
+	return z;
+}
+
 std::vector<int> kmp(const std::string &str, const std::string &pattern) {
 	std::vector<int> pi = prefix_function(pattern + "#" + str);
 	std::vector<int> result;
diff --git a/string_test.cpp b/string_test.cpp
index c51f548..bdc43e5 100644
--- a/string_test.cpp
+++ b/string_test.cpp
@@ -14,6 +14,14 @@ TEST_CASE("prefix function behaves as expected", "[string]") {
 	REQUIRE(prefix_function(v) == std::vector<int>{0, 1, 0, 0, 0, 0, 1, 2, 3});
 }
 
+TEST_CASE("z function behaves as expected", "[string]") {
+	REQUIRE(z_function("baobaba") == std::vector<int>{0, 0, 0, 2, 0, 2, 0, 0});
+	std::string s = "baobaba";
+	REQUIRE(z_function(s) == std::vector<int>{0, 0, 0, 2, 0, 2, 0});
+	std::vector<int> v = {2, 2, 8, 1, 5, 3, 2, 2, 8};
+	REQUIRE(z_function(v) == std::vector<int>{0, 1, 0, 0, 0, 0, 3, 1, 0});
+}
+
 TEST_CASE("kmp behaves as expected", "[string]") {
 	std::string s = "ababa";
 	std::string p = "aba";