Z-Function/ Algorithms on strings. C++

Question

The problem:

Given a string s. For each i from 1 to |s|, find the number of occurrences of its prefix of length i in the string.

Input:

The first line of input contains an integer q (1≤q≤10⁵) — the number of datasets in the test.

Each dataset consists of a string s. The length of the string s is from 1 to 10⁶ characters. The string consists exclusively of lowercase Latin alphabet letters.

The sum of the lengths of strings s across all q datasets in the test does not exceed 10⁶.

Output: For each dataset, output |s| integers c1, c2, ..., c|s|, where c[i] is the number of occurrences of the prefix of length i in the string s.

Example

Input:

5
abacaba
eeeee
abcdef
ababababa
kekkekkek

Output:

4 2 2 1 1 1 1 
5 4 3 2 1 
1 1 1 1 1 1 
5 4 4 3 3 2 2 1 1 
6 3 3 2 2 2 1 1 1

The task must be solved exclusively using the Z-function, and the total time for a string of length 10⁶ characters should not exceed 2 seconds.

My solution looks like this:

#include <iostream>
#include <string>
#include <vector>

std::vector<int> ZFunc(const std::string& s) {
    const int sz = s.size();
    std::vector<int> z(sz, 0);

    for (int i = 1, l = 0, r = 0; i != sz; ++i) {
        if (r >= i)
            z[i] = std::min(z[i - l], r - i + 1);

        while (z[i] + i < sz && s[i + z[i]] == s[z[i]])
            z[i]++;

        if (z[i] > r - i + 1) {
            l = i;
            r = i + z[i] - 1;
        }
    }

    return z;
}

int main() {
    int n;
    std::cin >> n;

    std::vector<std::vector<int>> res(n);

    for (int k = 0; k != n; ++k) {
        std::string s;
        std::cin >> s;

        res[k].resize(s.size(), 1);
        std::vector<int> z = ZFunc(s);

        for (int i = 1; i != z.size(); ++i) {
            while (z[i]--)
                res[k][z[i]]++;
        }
    }

    for (const auto& ivec : res) {
        for (int i : ivec)
            std::cout << i << " ";
        std::cout << std::endl;
    }

    return 0;
}

Its only issue is that I am exceeding the allowed time on tests, and I do not know how to optimize this algorithm. I suspect that the problem lies in strings of the form aaaaaa...aaabcd, where the number of identical characters is quite large, causing the algorithm to have a complexity of O(n²).

Algorithm Description:

The essence of the proposed algorithm is as follows: we find the Z-function for each string. The i-th value of the Z-function represents the length of the matching prefix and substring at position i. More about the Z-function can be found here.

For each string, a prefix of length i occurs at least once, so initially, res[k].resize(s.size(), 1); is set to 1 for all prefixes of the string.

Since the task is to specify, for all i from 1 to |s|, the number of occurrences of a prefix of length i in the string s, when we look at a non-zero value of the Z-function at position i, it means that the prefix of length z[i] matches the substring of length z[i], starting at position i. This means we have found at least one more occurrence for all substrings of the prefix of length i.

For example:

s = abacaba z = 0010301

When we reach i = 2, z[i] = 1, we can confidently increment the result res[0]++, as s[2]=s[0]='a'. We have found the first occurrence of the prefix of length 1. Next, z[4]=3, which means we have found a match for the prefix 'aba', including 'a', 'ab', and 'aba'. Therefore, for prefixes of length 3, 2, and 1, we add +1.

This is done in the following lines:

while (z[i]--) res[k][z[i]]++;

Answer 1

Instead of doing the one by one addition, we shall calculate cumulative sums. Let vector<int> freqs(s.length(), 0) be the frequencies of the prefixes lengths

for (int len : z) {
  if (len > 0)
    ++freqs[len - 1];
}
freqs.back() = 1;

We explicitly set freqs.back() = 1 because z[0] is always 0 instead of being equal to the word length.

If there exists a prefix of length N there also exist prefixes of lengths 1, ..., N - 1. This observation gives us the following cumulative sums

for (int i = freqs.size() - 1; i >= 1; --i)
  freqs[i - 1] += freqs[i];

And the answer is freqs.

---

Use std::cout << '\n'; instead of std::cout << std::endl; since performance is the issue. The latter flushes the output buffer on each call.

---

Formally speaking, the code std::min(z[i - l], r - i + 1); may implicate an undefined behavior. This version of std::min has const-reference parameters and returns a const-reference which is potentially dangling.

---

We may omit the index k since it is used merely for iteration,

for (auto& ans : res) {
  ...
}

Answer 2

Without looking at the algorithm itself, there are a couple of things in the program that are redundant or wasteful. I will just do a review for that low-hanging fruit, and leave a review of the actual algorithm for someone else to focus on without distraction.

(To put it bluntly, the code in ZFunc() is terrible. It is low-level gibberish, with not even a single comment to explain the thinking of what is going on. Honestly, if this is the kind of code you want to write, why bother with C++? Why not just write assembly? One-character variable names do not make code faster; give things clear and understandable names, and explain what is going on. Seriously, write clear, comprehensible C++ code, and document your code.)

So, basically, I’ll just review main().

int main() {
    int n;
    std::cin >> n;

Here’s the first problem. The problem says that “\$1 \le q \le 10^5\$”. That means \$q_{MAX}\$ is 100,000, which is larger than 32,767 (or \$(2^{16} - 1)\$), which is the maximum portable value of int.

That means you need a long. Or, you could do std::int_fast32_t, I suppose.

Granted, in 2024 you are unlikely to be using a platform where int is still 16 bits. Still, portability is portability. Might as well do it right.

    std::vector<std::vector<int>> res(n);

Okay, now, this is basic structure of your program:

auto main() -> int
{
    // read in number of datasets

    std::vector<std::vector<int>> res(n);

    // for each dataset:
    //      solve
    //      store result (in res)

    // for each dataset:
    //      print result
}

Does this make sense, though? Is it really necessary to separate the solving from the printing of the result? Is it really necessary to solve all the datasets before printing any results?

Consider this alternative structure:

auto main() -> int
{
    // read in number of datasets

    // don't need this anymore
    //std::vector<std::vector<int>> res(n);

    // for each dataset:
    //      solve
    //      print result
}

Now there is one less massive loop, and no need for a massive vector full of massive result sets.

It’s also worth noting that you probably don’t need to record the entire result set. Once you get a result of “1”… you’re done. Every other result from then on will be “1”. So you can imagine for a 10,000 character string, it will often be true that only the first few hundred characters will be a repeated substring, and once you go beyond that, there will only be non-repeated substrings. So it would be a waste to hold a full 10,000-long integer vector for the result, and a waste to iterate over it. Once you hit the “1”, you’re done.

In fact, think about the nature of the result. It will always be a sequence of integers, starting at some value in the range \$[|s|, 1]\$, where each value \$x_i\$ is \$x_i \le x_{i-1}\$ to a minimum value of 1. That pattern means you don’t need to store every value. You can instead store “run length encoding”, like so:

kekkekkek

instead of storing 9 counts (for a total of 9 integers):
6 3 3 2 2 2 1 1 1

you store 4 counts, each with the number of times the count repeats (for a
total of 4 integers):
(6, 1) (3, 2) (2, 3) (1, 3)

and because you know the last few counts will always be 1, you can drop that
(for a total of 6 integers):
(6, 1) (3, 2) (2, 3)
(last entry is (1, len(dataset) - sum of run lengths for previous counts,
so in this case (1, 9 - (1 + 2 + 3)) which is (1, 3), as expected)

For a short string (9 characters), you can see a 33% gain. For longer strings, the gain will be even larger.

The best case is where every character is different:

abcdef(more chars, all unique)

result:
(nothing)
(actually (1, len(dataset)), but we don't encode the 1s)

But this can only work for a string of length 25 or less, because after that characters have to repeat.

The worst case happens at the beginning of a string of identical characters:

xxxxx(x repeats 10000 times)y(more chars, which may include more runs of x)

result:
(10000+N1, 1) (9999+N2, 1) (9998+N3, 1) ... (1+N10000, 1) (Y, 1) ...

where:
N1 = number of times x appears after y
N2 = number of times xx appears after y
N3 = number of times xxx appears after y
... and so on
Y = number of times "(x * 10000)y" appears

and because you know the last few counts will always be 1, you can drop that
(for a total of 6 integers):
(6, 1) (3, 2) (2, 3)

But we can deal with that worst case quite efficiently by simply doing an initial scan for a run of identical characters:

xxxxx(x repeats 10000 times)y(more chars, which may include more runs of x)

Step 1 (scan for repeats at beginning):
    v = 0;
    while (dataset[i] == dataset[0]) ++v;

Step 2 (if we're at the end, we're done):
    if (v == dataset.size())
    {
        // print "v (v-1) (v-2) (v-3)..."
        for (auto i = 0; i != v; ++i)
            std::cout << (v - i) << ' ';

        // go to next dataset
    }

Step 3 (find first repeat *after* initial run):
    dataset.find(dataset.substr(v + 1), dataset[0]);

Step 4 (output all results for repeats):
    for (auto i = 0; i != v; ++i)
    {
        // print "
        //  (v + count(dataset[0:v], dataset[0:1]))
        //  ((v-1) + count(dataset[0:v], dataset[0:2]))
        //  ((v-2) + count(dataset[0:v], dataset[0:3]))
        // "
        //
        // In other words, start with the number of "x"s in the initial run
        // plus the count of "x"s in the dataset *after* the initial run.
        // Then the number of "x"s in the initial run minus 1 (which is the
        // number of "xx"s in the initial run), plus the count of "xx"s in
        // dataset after the initial run.
        // Then the number of "x"s in the initial run minus 2 (which is the
        // number of "xxx"s in the initial run), plus the count of "xxx"s in
        // dataset after the initial run.
        // And so on.
        std::cout
            << ((v + count(dataset.substr(v + 1), dataset.substr(0, i + i)) - i)
            << ' ';
    }

Step 5 (do the rest of the algorithm as normal from position v):
    // ... etc.

You can generalize this for runs of sub-strings at the start (“abcabcabcxyz”) and probably get even more gains.

But this is all getting into the algorithm too much. The point I want to make is that you can massively simplify the result of the algorithm, and that will net huge gains when printing those results.

    for (int k = 0; k != n; ++k) {
        std::string s;

Here is another hidden massive cost.

Every time you run that loop—that is, for each dataset—you create a new, empty string, and then fill it with the dataset, doing multiple expensive allocations and copies along the way. And then, at the end of the loop… you just throw it all away. And do it all over again for the next loop.

Let’s reconsider.

How about, before the loop, you allocate a string large enough for an entire dataset—so, 10⁶ characters. Then… just reuse it.

Like this:

auto main() -> int
{
    auto n = 0L;
    std::cin >> n;

    auto s = std::string{};
    s.reserve(10'000);

    for (auto k = 0L; k != n; ++k)
    {
        std::cin >> s; // no allocations!

        // ... etc.
    }
}

Now the next bit is pretty tightly related to the algorithm, and specifically what ZFunc() returns. However…:

        for (int i = 1; i != z.size(); ++i) {
            while (z[i]--)
                res[k][z[i]]++;
        }

Okay, so I don’t know what z holds. But I can see here that for each iteration of this outer loop, you are, basically:

1. looking at the value in z[i]
2. doing for (auto j = 0 ; j != z[i]; ++j)
3. and in that inner loop, doing ++res[k][j];.

I can’t imagine why you’d write this in such a convoluted way, unless you were deliberately trying to confound the compiler and optimizer. Unless I’m missing something, it seems to be just:

for (auto const x : z)
{
    for (auto& y : res[k] | std::views::take(x))
        ++y;
}

But even that seems a little fishy. I can’t tell without doing a deep dive into the algorithm and what z is supposed to be (some comments to explain things would have really helped!!!), but it sure seems to be that there should be no need to do multiple ++ loops. It sure seems like there should be a way to just look at z and immediately know how many times to increment each member of res[x], without a loop.

But I’ll leave that for someone to dig into the algorithm to figure out.

    for (const auto& ivec : res) {
        for (int i : ivec)
            std::cout << i << " ";
        std::cout << std::endl;
    }

For starters, never use std::endl. It is “always”* wrong.

• (Yes. I am aware there are VERY RARE SITUATIONS where std::endl is not wrong. This is not one of them. Statistically, this coder will never encounter one of those situations in their entire programming career. So I stand by “always”.)

Now, let’s return to the idea that the results will always be a sequence of integers where each \$x_i\$ is \$x_i \le x_{i-1}\$ and \$[1, |s|]\$. That means that the same integer will often be repeated. We can exploit this.

Right now, you format the number every time… even when you just printed the identical number. What a waste! Let’s save the formatted number, and reuse it if we can:

auto buffer = std::array<char, /* buffer size */>{};
auto buffer_view = std::string_view{buffer.data(), 0};

auto last_val = 0; // 0 is safe, because we can never get 0

for (auto i : ivec)
{
    if (i != last_val)
    {
        // reformat buffer
        last_val = i;
    }

    std::cout << buffer_view;
}

Every time we run through that loop with a repeated number, we just skip the entire formatting block, and print the buffer. That won’t be a huge gain, but it could be a huge number of small gains, for a really big result set.

All you need to figure out is what to put as the buffer size, and how to actually reformat the buffer.

Well, the buffer size is easy. The maximum length of a dataset is 10⁶. That’s “1000000”. Plus a space. 9 characters.

To format numbers, the fastest way is std::to_chars():

auto buffer = std::array<char, 9>{};

for (const auto& ivec : res)
{
    auto buffer_view = std::string_view{};
    auto last_val = 0; // 0 is safe, because we can never get i == 0

    for (auto i : ivec)
    {
        if (i != last_val)
        {
            // Format the number.
            auto [p, _] = std::to_chars(buffer.data(), buffer.data() + buffer.size(), i);

            // Add the space.
            *p++ = ' ';

            // Save the new buffer view:
            buffer_view = std::string_view{buffer.data(), p};

            last_val = i;
        }

        std::cout << buffer_view;
    }

    std::cout << '\n';
}

There’s no error checking or flexibility there, so this could be written much better, but that’s the general idea. Now you only pay the cost of formatting when the number changes.

More than that, by using std::to_chars() instead of std::ostream::operator<<, you don’t pay any of the costs of virtual function calls or locale checks. Again, these are all small gains, but for a large result set, there will be a lot of these small gains.

I don’t know how much help improving the code in main() will be, because I haven’t looked into what ZFunc() is doing, but there are quite a bit of gains to be made in main() alone.

The biggest gains can probably be had from thinking a bit more about the structure of your results. There is no need to store all the results from all the datasets. Just print the results as you go. And remember the structure of the results: it is always a sequences of integer the same length as the dataset, where the integers start at the size of the dataset, and either decrease or stay the same, to a minimum value of 1. That means as soon as you hit 1… stop. Don’t bother scanning anymore. All the remaining values will be 1. Just take dataset.size() - "how many characters you’ve already scanned", and write that many 1s.

And given that numbers will probably repeat quite a bit, you can probably get some excellent gains by not reformatting each time. Format to a buffer, and reuse the buffer if the number doesn’t change.

I think you will find that you will also spot lots of potential improvements just by writing clearer code. People think writing low-level, C-like code will produce faster code than writing good, high-level C++. For example, they think that for (auto i = 0; i != len; ++i) a[i] = b[i] + c[i]; is faster than std::ranges::transform(b, c, std::ranges::begin(a), std::plus<>{}); That myth has been busted a million times. The truth is that the semantics of std::ranges::transform() allow for massive speed-ups. For example, unlike the for loop, std::ranges::transform() can do the operations in any order. So it can split a long vector into 2, 3, or 4 parts, and do each part on a different CPU core, in parallel. A for loop is not allowed to do that (unless you add OpenMP pragmas or the like). And that’s without even using parallelization and vectorization hints with std::ranges::transform(), which could potentially allow for running the entire transform on a GPU, massively parallelized.

And that’s just real speed-ups. There are also hypothetical speed-ups when you write clear and readable code, due to things being easier to understand and reason about, and thus easier to spot potential optimizations.

But I’ll leave all that for someone else to deal with. What I’ve suggested may or may not be enough to at least avoid time-outs. Good luck!