cp-library
競技プログラミング用 C++ ライブラリ
Z algorithm
string/z_algorithm.hpp · GitHub で見る
共通ページ
#include "lib/string/z_algorithm.hpp"
auto z = z_algorithm("ababa"); // {5,0,3,0,1}
auto a = z_algorithm(std::vector<int>{1, 2, 1, 2}); // {4,0,2,0}
z_algorithm(sequence) は、各位置からの接尾辞と列全体との最長共通接頭辞の長さを vector<int> で返します。非空列では z[0]=sequence.size()、空列では空配列です。
string_view、vector<T>、span<const T> を渡せます。T は等値比較が可能とし、vector<bool> は対象外です。列長は INT_MAX 以下とします。
列長を \(n\) として、時間・追加空間は \(O(n)\) です。
依存ライブラリ
コード
#ifndef STRING_Z_ALGORITHM_HPP
#define STRING_Z_ALGORITHM_HPP
#include "../components/sequence_match.hpp"
#include <cassert>
#include <limits>
#include <span>
#include <string_view>
#include <vector>
template <class T>
std::vector<int> z_algorithm(std::span<const T> sequence) {
assert(sequence.size() <= std::size_t(std::numeric_limits<int>::max()));
std::vector<int> result(sequence.size());
sequence_match_detail::Matcher<T> match(sequence);
sequence_match_detail::z_table(int(sequence.size()), result.data(), [&](int a, int b, int limit) { return match.forward(a, b, limit); });
return result;
}
inline std::vector<int> z_algorithm(std::string_view sequence) {
return z_algorithm(std::span<const char>(sequence.data(), sequence.size()));
}
template <class T, class Allocator>
std::vector<int> z_algorithm(const std::vector<T, Allocator> &sequence) {
return z_algorithm(std::span<const T>(sequence.data(), sequence.size()));
}
#endif
#line 1 "string/z_algorithm.hpp"
#line 1 "components/sequence_match.hpp"
#line 1 "components/sequence_compare.hpp"
#include <bit>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <type_traits>
#if defined(__x86_64__)
#include <immintrin.h>
#endif
namespace sequence_compare_detail {
#if defined(__x86_64__)
__attribute__((target("avx2"))) inline std::size_t forward_avx2(const unsigned char *a, const unsigned char *b, std::size_t size) {
std::size_t i = 0;
for (; size - i >= 32; i += 32) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i));
auto difference = ~std::uint32_t(_mm256_movemask_epi8(_mm256_cmpeq_epi8(x, y)));
if (difference) return i + std::countr_zero(difference);
}
while (i < size && a[i] == b[i]) ++i;
return i;
}
__attribute__((target("avx2"))) inline std::size_t backward_avx2(const unsigned char *a, const unsigned char *b, std::size_t size) {
std::size_t i = 0;
for (; size - i >= 32; i += 32) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a - i - 32));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(b - i - 32));
auto difference = ~std::uint32_t(_mm256_movemask_epi8(_mm256_cmpeq_epi8(x, y)));
if (difference) return i + std::countl_zero(difference);
}
while (i < size && a[-std::ptrdiff_t(i) - 1] == b[-std::ptrdiff_t(i) - 1]) ++i;
return i;
}
#endif
template<bool Vectorize = true, class T>
inline int prefix(const T *first, const T *second, int limit) {
if (!limit) return 0;
int matched = 0;
if constexpr (std::is_integral_v<T> && std::has_unique_object_representations_v<T>) {
const auto *a = reinterpret_cast<const unsigned char *>(first);
const auto *b = reinterpret_cast<const unsigned char *>(second);
std::size_t offset = 0, size = std::size_t(limit) * sizeof(T);
while (size - offset >= 8) {
std::uint64_t x, y;
std::memcpy(&x, a + offset, 8);
std::memcpy(&y, b + offset, 8);
auto difference = x ^ y;
if (difference) {
int bits = std::endian::native == std::endian::little ? std::countr_zero(difference) : std::countl_zero(difference);
return static_cast<int>((offset + bits / 8) / sizeof(T));
}
offset += 8;
#if defined(__x86_64__)
if constexpr (Vectorize) if (size - offset >= 64 && __builtin_cpu_supports("avx2"))
return static_cast<int>((offset + forward_avx2(a + offset, b + offset, size - offset)) / sizeof(T));
#endif
}
matched = static_cast<int>(offset / sizeof(T));
}
while (matched < limit && first[matched] == second[matched]) ++matched;
return matched;
}
template<bool Vectorize = true, class T>
inline int suffix(const T *first_end, const T *second_end, int limit) {
if (!limit) return 0;
int matched = 0;
if constexpr (std::is_integral_v<T> && std::has_unique_object_representations_v<T>) {
const auto *a = reinterpret_cast<const unsigned char *>(first_end);
const auto *b = reinterpret_cast<const unsigned char *>(second_end);
std::size_t offset = 0, size = std::size_t(limit) * sizeof(T);
while (size - offset >= 8) {
std::uint64_t x, y;
std::memcpy(&x, a - offset - 8, 8);
std::memcpy(&y, b - offset - 8, 8);
auto difference = x ^ y;
if (difference) {
int bits = std::endian::native == std::endian::little ? std::countl_zero(difference) : std::countr_zero(difference);
return static_cast<int>((offset + bits / 8) / sizeof(T));
}
offset += 8;
#if defined(__x86_64__)
if constexpr (Vectorize) if (size - offset >= 64 && __builtin_cpu_supports("avx2"))
return static_cast<int>((offset + backward_avx2(a - offset, b - offset, size - offset)) / sizeof(T));
#endif
}
matched = static_cast<int>(offset / sizeof(T));
}
while (matched < limit && first_end[-matched - 1] == second_end[-matched - 1]) ++matched;
return matched;
}
}
#line 1 "components/z_algorithm.hpp"
namespace sequence_match_detail {
template <class F>
void z_table(int n, int *z, F match) {
if (!n) return;
z[0] = n;
for (int i = 1, base = 0, end = 0; i < n; ++i) {
int length = 0;
if (i < end) {
length = z[i - base];
if (length < end - i) {
z[i] = length;
continue;
}
length = end - i;
}
length += match(length, i + length, n - i - length);
z[i] = length;
if (i + length > end) base = i, end = i + length;
}
}
template <class F, class Emit>
void prefix_matches(int n, int m, const int *z, F match, Emit emit) {
for (int i = 0, base = 0, end = 0; i < n; ++i) {
int length = 0;
if (i < end) {
length = z[i - base];
if (length < end - i) {
emit(i, length);
continue;
}
length = end - i;
}
length += match(length, i + length, m - length);
if (i + length > end) base = i, end = i + length;
emit(i, length);
}
}
}
#line 6 "components/sequence_match.hpp"
#include <algorithm>
#include <array>
#include <span>
#include <utility>
#include <vector>
namespace sequence_match_detail {
#if defined(__x86_64__)
using sequence_compare_detail::forward_avx2, sequence_compare_detail::backward_avx2;
#endif
template <class T, bool Vectorize = true>
class Matcher {
const T *data;
std::vector<std::array<int, 2>> runs;
public:
explicit Matcher(std::span<const T> values) : data(values.data()) {}
void index_runs(int n) {
runs.resize(n);
for (int first = 0, last; first < n; first = last) {
last = first + 1;
while (last < n && data[first] == data[last]) ++last;
for (int i = first; i < last; ++i) runs[i] = {first, last};
}
}
std::pair<int, int> forward_counted(int i, int j, int limit) const {
if (runs.empty() || !limit || data[i] != data[j]) {
int length = forward(i, j, limit);
return {length, length};
}
int a = runs[i][1] - i, b = runs[j][1] - j, skip = std::min({limit, a, b});
if (a != b || skip == limit) return {skip, 1};
int length = forward(i + skip, j + skip, limit - skip);
return {skip + length, length + 1};
}
std::pair<int, int> backward_counted(int i, int j, int limit) const {
if (runs.empty() || !limit || data[i - 1] != data[j - 1]) {
int length = backward(i, j, limit);
return {length, length};
}
int a = i - runs[i - 1][0], b = j - runs[j - 1][0], skip = std::min({limit, a, b});
if (a != b || skip == limit) return {skip, 1};
int length = backward(i - skip, j - skip, limit - skip);
return {skip + length, length + 1};
}
int forward(int i, int j, int limit) const {
if (!limit) return 0;
return sequence_compare_detail::prefix<Vectorize>(data + i, data + j, limit);
}
int backward(int i, int j, int limit) const {
if (!limit) return 0;
return sequence_compare_detail::suffix<Vectorize>(data + i, data + j, limit);
}
};
}
#line 5 "string/z_algorithm.hpp"
#include <cassert>
#include <limits>
#line 8 "string/z_algorithm.hpp"
#include <string_view>
#line 10 "string/z_algorithm.hpp"
template <class T>
std::vector<int> z_algorithm(std::span<const T> sequence) {
assert(sequence.size() <= std::size_t(std::numeric_limits<int>::max()));
std::vector<int> result(sequence.size());
sequence_match_detail::Matcher<T> match(sequence);
sequence_match_detail::z_table(int(sequence.size()), result.data(), [&](int a, int b, int limit) { return match.forward(a, b, limit); });
return result;
}
inline std::vector<int> z_algorithm(std::string_view sequence) {
return z_algorithm(std::span<const char>(sequence.data(), sequence.size()));
}
template <class T, class Allocator>
std::vector<int> z_algorithm(const std::vector<T, Allocator> &sequence) {
return z_algorithm(std::span<const T>(sequence.data(), sequence.size()));
}