EXPREP - Editorial

PROBLEM LINK:

Practice
Div-2 Contest
Div-1 Contest

Author: aviroop123
Tester: enchom
Editorialist: aviroop123

DIFFICULTY:

HARD.

PREREQUISITES:

Suffix Automata, DSU on Tree trick, Binary Indexed Tree, Suffix Array, Sparse Segment Tree.

PROBLEM:

Find the sum of powers of all substrings of a given string S. The power of a string T is defined as the sum of weights of all its periodic substrings. A string R is said to be a periodic substring of string T if 1 \leq |R| \leq |T| and T = R + R + ... + R + P where P is a prefix of R and + denotes concatenation.

QUICK EXPLANATION:

  • Let’s reverse S.
  • if there are two equal substrings S[a : b] and S[c : d] such that. b \lt d, then S[b + 1 : d] is a periodic substring of S[a: d].
  • Let’s say a substring T ends at positions i_1, i_2, ..., i_K in the string S. Then the contribution of T to the total sum is \sum_{p = 1}^{k}\sum_{q = p+1}^{k} W(S[i_{p} : i_{q}]). To calculate this, we use suffix automata and maintain the endpos class of each node of the automata using DSU on the link tree and also maintain two BITs for storing prefix sums.

Complexity : O(N \cdot log^2N).

EXPLANATION:

Subtask 1

Iterate over all the substrings and all the lengths of the repeating substrings. The repeating substring will always be a prefix of the substring considered. Then Just check whether the substring is of the given form or not and add the weight accordingly.

Complexity : O(N^3) or O(N^3 \cdot logN) or O(N^4) according to the implementation.

Subtask 2

First let’s reverse the given string S so that for a particular substring, instead of r being a prefix, r becomes a suffix.

Now, if there are two equal substrings S[a : b] = S[c : d] and b \lt d, then S[b + 1 : d] is a periodic suffix of S[a : d]. Moreover, there is a bijection between all such pairs of equal substrings ([a : b], [c : d]) and (T, R) where |R| \lt |T|. The case |R| = |T| can be handled exclusively.

As an example, consider string S = aabaaba.
Since S[1 : 4] = S[4 : 7] = aaba are two equal substrings, we can see that R = S[5 : 7] = aba is a periodic suffix of T = S[1 : 7] = aabaaba.

The ending positions of the substrings can be maintained for each distinct substrings of the original string and the corresponding weight can be calculated using prefix sums. Hashing or any suffix structure might do the trick.

Complexity : O(N^2) or O(N^2 \cdot logN) depending on the implementation.

Subtask 3

First we’ll describe a bit about suffix automata and what is an endpos class of a substring, then we’ll reduce it to a problem related to trees and finally I’ll discuss how to compute it efficiently using the DSU on tree trick and BITs.

Suffix Automata Algorithm Details

Suffix automata basically breaks down each substring of a string and groups them into endpos classes. Substrings with the same endpos class are grouped together. It can be proven that number of such groups is O(N). For more details about suffix automata, you can refer to this awesome blog on Suffix Automata.

What is an endpos class in Suffix Automata ?

An endpos class of a substring T present in S Is defined as the set of all ending positions of S where T occurs as a substring. For example, if T = "ab" and S = "ababacbab", then endpos(T) = \{2, 4, 9\}.

Let pre[i] denote weight of prefix of S till i. For each distinct substring T of S, let endpos(T) = \{i_1, i_2, ..., i_k \}. The contribution of T to the total answer would be \sum_{p = 1}^{k}\sum_{q=p+1}^{k}W(S[i_p : i_q]) = \sum_{p = 1}^{k}\sum_{q=p+1}^{k}(pre[i_q] - pre[i_p]).

Simplifying the summation we get \sum_{q=1}^{k}pre[i_q] \cdot (q - 1) - \sum_{p=1}^{k}pre[i_p] \cdot (k-p).

Note that the contribution of each string having the same endpos class is the same! And we know that the number of such endpos classes is linear.

Properties of Suffix Automata

Some properties of the automata which can help to compute this sum efficiently are:

  1. The links present in the suffix automata determines a rooted tree. The root of the tree is the node corresponding to the empty substring.
  2. For two nodes A and B that are not ancestors of one another, endpos(A) \cap endpos(B) = \phi.
  3. We define terminal states as those nodes which have a prefix of S ending at them. For each terminal node A present in the suffix automata, we define Z(A) as the index i s.t. the prefix ending at index i terminates at node A of the suffix automata. Z(A) is null for non-terminal nodes. Then for a node A, endpos(A) = endpos(C_1) \cup endpos(C_2) \cup ... endpos(C_k) \cup Z(A) .

So the problem is reduced to a tree problem where some nodes have an index assigned to them ( the terminal nodes) and we compute the summation for each node of the tree. endpos classes and the required sums can be maintained using DSU on Tree trick and two Binary Indexed Trees.

Refer to the setter code for more details.

Complexity : O(N log^2 N)

ALTERNATE EXPLANATION:

The problem can be also solved with suffix arrays.

The idea is that after building the suffix array, we merge ranges in decreasing order of length and sum them.

First we create the suffix array and the lcp array of the given string. Then, we merge suffix in the suffix array by sorting all the lcp values and connecting the suffixes whose lcp value is greater than the current length. Merging can be done by maintaining dsu on an array.

Now, for a particular length, we have some connected components. For each component, we maintain a sparse segment tree. When connecting two components, we can use a dfs to merge the two segtrees, since the merge would require O((smaller segtree size) * logN).

In each node of the segtree we maintain two values, 1) the number of suffixes in the range 2) the sum of prefix sums in the range. Here range denotes the range that the node is covering in the segtree.

Refer to the code and the comments for more details regarding this approach.

Complexity : O(N log^2N)

SOLUTIONS:

Setter's Solution Using Suffix Automata
#include "bits/stdc++.h"
using namespace std;
#ifndef LOCAL
#define endl '\n'
#endif

#define fr(i, a, b) for(int i = a; i <= b; i++)
#define rf(i, a, b) for(int i = a; i >= b; i--)
#define pb push_back
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define sz(x) (int)x.size()

typedef long long ll;

void sc() {}

template <typename Head, typename... Tail>
void sc(Head &H, Tail &... T) { cin >> H; sc(T...); }

#ifdef LOCAL
#define debug(...) cerr << "[L:" << __LINE__ << "][" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

const int mod = 998244353;

int pwr(int a,ll b) {
    int ans = 1;
    while(b) {
        if(b & 1) ans = (ans * 1LL * a) % mod;
        a = (a * 1LL * a) % mod;
        b >>= 1;
    }
    return ans;
}

struct state {
    int len, link;
    map<char, int> next;
};
const int M = 1e6 + 5; // 2 * maxlen

int Z[M];

state st[M];
int sz, last;
void sa_init() {
    while(sz > 0) {
        sz--;
        st[sz].next.clear();
        Z[sz] = 0;
    }
    st[0].len = 0;
    st[0].link = -1;
    sz++;
    last = 0;
}
void sa_extend(char c, int ck) {
    // ck -> index of character c to be extended
    int cur = sz++;
    st[cur].len = st[last].len + 1;
    int p = last;
    while (p != -1 && !st[p].next.count(c)) {
        st[p].next[c] = cur;
        p = st[p].link;
    }
    if (p == -1) {
        st[cur].link = 0;
    } else {
        int q = st[p].next[c];
        if (st[p].len + 1 == st[q].len) {
            st[cur].link = q;
        } else {
            int clone = sz++;
            st[clone].len = st[p].len + 1;
            st[clone].next = st[q].next;
            st[clone].link = st[q].link;
            while (p != -1 && st[p].next[c] == q) {
                st[p].next[c] = clone;
                p = st[p].link;
            }
            st[q].link = st[cur].link = clone;
        }
    }
    last = cur;
    Z[cur] = ck; // index of prefix ending at cur
}

vector<int> g[M];
int S[M], str[M], en[M], ver[M], tt;

void pre_dfs(int u,int p) {
    S[u] = 1;
    str[u] = ++tt;
    ver[tt] = u;
    for(int v : g[u]) {
        if(v != p) {
            pre_dfs(v, u);
            S[u] += S[v];
        }
    }
    en[u] = tt;
}

int w[26];
int pre[M];

template<typename T, size_t N>
struct bit {
    int BIT[N];
    int n;
    void init(int nn) {
        n = nn + 2;
        fr(i, 0, n) BIT[i] = 0;
    }
    void upd(int idx, int val) {
        idx++;
        if(val < 0) val += mod;
        while(idx < n) {
            BIT[idx] += val;
            if(BIT[idx] >= mod) BIT[idx] -= mod;
            idx += idx & (-idx);
        }
    }
    int query(int idx) {
        idx++;
        int s = 0;
        while(idx > 0) {
            s += BIT[idx];
            if(s >= mod) s -= mod;
            idx -= idx & (-idx);
        }
        return s;
    }
    int query(int l,int r) {
        int q = (query(r) - query(l - 1) + mod);
        if(q >= mod) q -= mod;
        return q;
    }
};

bit<int, M / 2 + 2> b1, b2;
int n;

int ans, cur;
void add(int u,int x) {
    // q -> contribution of u to the current sum
    int q = (b2.query(u, n) - b1.query(u, n) * 1LL * pre[u]) % mod;
    q += (b1.query(0, u - 1) * 1LL * pre[u] - b2.query(0, u - 1)) % mod;
    if(q >= mod) q -= mod;
    if(q < 0) q += mod;
    cur += x * q;
    if(cur < 0) cur += mod;
    if(cur >= mod) cur -= mod;
    // update the BITs
    b1.upd(u, x);
    b2.upd(u, x * pre[u]);
}

void dfs(int v, int p, bool keep) { // dsu on tree trick
    int mx = -1, bigChild = -1;
    for(auto u : g[v])
    if(S[u] > mx)
        mx = S[u], bigChild = u;
    for(auto u : g[v])
        if(u != bigChild)
            dfs(u, v, 0);
    if(bigChild != -1) dfs(bigChild, v, 1);
    for(auto u : g[v])
        if(u != bigChild)
            for(int p = str[u]; p <= en[u]; p++)
                if(Z[ver[p]])
                    add(Z[ver[p]], 1);
    if(Z[v])
        add(Z[v], 1);
    if(v) {
        ans = (ans + (st[v].len - st[st[v].link].len) * 1LL * cur) % mod;
    }
    if(ans >= mod) ans -= mod;
    if(!keep)
        for(int p = str[v]; p <= en[v]; p++)
            if(Z[ver[p]])
                add(Z[ver[p]], -1);
}

void solve() {
    ans = 0, cur = 0, tt = 0;
    string s;
    cin >> s;
    n = sz(s);
    b1.init(n), b2.init(n); // initializing binary indexed trees
    fr(i, 0, 25) {
        sc(w[i]);
        if(w[i] >= mod) w[i] -= mod;
    }
    fr(i, 1, n) {
        pre[i] = pre[i - 1] + w[s[i - 1] - 'a']; // presum of weights of characters of the string
        if(pre[i] >= mod) pre[i] -= mod;
    }
    sa_init();
    fr(i, 0, sz(s) - 1)
        sa_extend(s[i], i + 1);
    fr(i, 0, sz)
        g[i].clear();
    fr(i, 1, sz)
        g[st[i].link].pb(i); // building the tree
    pre_dfs(0, 0);
    dfs(0, 0, 0);
    int sum = 0;
    fr(i, 1, n) { // for handling the case where r is the whole substring.
        ans = (ans + pre[i] * 1LL * i - sum) % mod;
        if(ans < 0) ans += mod;
        sum += pre[i];
        if(sum >= mod) sum -= mod;
    }
    int den = pwr((n * 1LL * (n + 1)) % mod, mod - 2);
    ans = (ans * 2LL * den) % mod;
    cout << ans << endl;
}

int main() {
    ios :: sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    int t = 1;
    cin >> t;
    for(int tt = 1; tt <= t; tt++) {
        solve();
    }
    return 0;
}
Setter's Solution Using Suffix Array
#pragma GCC optimize("Ofast")
#include "bits/stdc++.h"
using namespace std;
#ifndef LOCAL
#define endl '\n'
#endif

#define fr(i, a, b) for(int i = a; i <= b; i++)
#define rf(i, a, b) for(int i = a; i >= b; i--)
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()

typedef long long ll;

void sc() {}

template <typename Head, typename... Tail>
void sc(Head &H, Tail &... T) { cin >> H; sc(T...); }

#ifdef LOCAL
#define debug(...) cerr << "[L:" << __LINE__ << "][" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

// -- Mint template (https://github.com/the-tourist/algo/blob/master/numeric/mint.cpp)

template <typename T>
T inverse(T a, T m) {
T u = 0, v = 1;
while (a != 0) {
    T t = m / a;
    m -= t * a; swap(a, m);
    u -= t * v; swap(u, v);
}
assert(m == 1);
return u;
}

template <typename T>
class Modular {
public:
using Type = typename decay<decltype(T::value)>::type;

constexpr Modular() : value() {}
template <typename U>
Modular(const U& x) {
    value = normalize(x);
}

template <typename U>
static Type normalize(const U& x) {
    Type v;
    if (-mod() <= x && x < mod()) v = static_cast<Type>(x);
    else v = static_cast<Type>(x % mod());
    if (v < 0) v += mod();
    return v;
}

const Type& operator()() const { return value; }
template <typename U>
explicit operator U() const { return static_cast<U>(value); }
constexpr static Type mod() { return T::value; }

Modular& operator+=(const Modular& other) { if ((value += other.value) >= mod()) value -= mod(); return *this; }
Modular& operator-=(const Modular& other) { if ((value -= other.value) < 0) value += mod(); return *this; }
template <typename U> Modular& operator+=(const U& other) { return *this += Modular(other); }
template <typename U> Modular& operator-=(const U& other) { return *this -= Modular(other); }
Modular& operator++() { return *this += 1; }
Modular& operator--() { return *this -= 1; }
Modular operator++(int) { Modular result(*this); *this += 1; return result; }
Modular operator--(int) { Modular result(*this); *this -= 1; return result; }
Modular operator-() const { return Modular(-value); }

template <typename U = T>
typename enable_if<is_same<typename Modular<U>::Type, int>::value, Modular>::type& operator*=(const Modular& rhs) {
#ifdef _WIN32
    uint64_t x = static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value);
    uint32_t xh = static_cast<uint32_t>(x >> 32), xl = static_cast<uint32_t>(x), d, m;
    asm(
    "divl %4; \n\t"
    : "=a" (d), "=d" (m)
    : "d" (xh), "a" (xl), "r" (mod())
    );
    value = m;
#else
    value = normalize(static_cast<int64_t>(value) * static_cast<int64_t>(rhs.value));
#endif
    return *this;
}
template <typename U = T>
typename enable_if<is_same<typename Modular<U>::Type, int64_t>::value, Modular>::type& operator*=(const Modular& rhs) {
    int64_t q = static_cast<int64_t>(static_cast<long double>(value) * rhs.value / mod());
    value = normalize(value * rhs.value - q * mod());
    return *this;
}
template <typename U = T>
typename enable_if<!is_integral<typename Modular<U>::Type>::value, Modular>::type& operator*=(const Modular& rhs) {
    value = normalize(value * rhs.value);
    return *this;
}

Modular& operator/=(const Modular& other) { return *this *= Modular(inverse(other.value, mod())); }

template <typename U>
friend const Modular<U>& abs(const Modular<U>& v) { return v; }

template <typename U>
friend bool operator==(const Modular<U>& lhs, const Modular<U>& rhs);

template <typename U>
friend bool operator<(const Modular<U>& lhs, const Modular<U>& rhs);

template <typename U>
friend std::istream& operator>>(std::istream& stream, Modular<U>& number);

private:
Type value;
};

template <typename T> bool operator==(const Modular<T>& lhs, const Modular<T>& rhs) { return lhs.value == rhs.value; }
template <typename T, typename U> bool operator==(const Modular<T>& lhs, U rhs) { return lhs == Modular<T>(rhs); }
template <typename T, typename U> bool operator==(U lhs, const Modular<T>& rhs) { return Modular<T>(lhs) == rhs; }

template <typename T> bool operator!=(const Modular<T>& lhs, const Modular<T>& rhs) { return !(lhs == rhs); }
template <typename T, typename U> bool operator!=(const Modular<T>& lhs, U rhs) { return !(lhs == rhs); }
template <typename T, typename U> bool operator!=(U lhs, const Modular<T>& rhs) { return !(lhs == rhs); }

template <typename T> bool operator<(const Modular<T>& lhs, const Modular<T>& rhs) { return lhs.value < rhs.value; }

template <typename T> Modular<T> operator+(const Modular<T>& lhs, const Modular<T>& rhs) { return Modular<T>(lhs) += rhs; }
template <typename T, typename U> Modular<T> operator+(const Modular<T>& lhs, U rhs) { return Modular<T>(lhs) += rhs; }
template <typename T, typename U> Modular<T> operator+(U lhs, const Modular<T>& rhs) { return Modular<T>(lhs) += rhs; }

template <typename T> Modular<T> operator-(const Modular<T>& lhs, const Modular<T>& rhs) { return Modular<T>(lhs) -= rhs; }
template <typename T, typename U> Modular<T> operator-(const Modular<T>& lhs, U rhs) { return Modular<T>(lhs) -= rhs; }
template <typename T, typename U> Modular<T> operator-(U lhs, const Modular<T>& rhs) { return Modular<T>(lhs) -= rhs; }

template <typename T> Modular<T> operator*(const Modular<T>& lhs, const Modular<T>& rhs) { return Modular<T>(lhs) *= rhs; }
template <typename T, typename U> Modular<T> operator*(const Modular<T>& lhs, U rhs) { return Modular<T>(lhs) *= rhs; }
template <typename T, typename U> Modular<T> operator*(U lhs, const Modular<T>& rhs) { return Modular<T>(lhs) *= rhs; }

template <typename T> Modular<T> operator/(const Modular<T>& lhs, const Modular<T>& rhs) { return Modular<T>(lhs) /= rhs; }
template <typename T, typename U> Modular<T> operator/(const Modular<T>& lhs, U rhs) { return Modular<T>(lhs) /= rhs; }
template <typename T, typename U> Modular<T> operator/(U lhs, const Modular<T>& rhs) { return Modular<T>(lhs) /= rhs; }

template<typename T, typename U>
Modular<T> power(const Modular<T>& a, const U& b) {
assert(b >= 0);
Modular<T> x = a, res = 1;
U p = b;
while (p > 0) {
    if (p & 1) res *= x;
    x *= x;
    p >>= 1;
}
return res;
}

template <typename T>
bool IsZero(const Modular<T>& number) {
return number() == 0;
}

template <typename T>
string to_string(const Modular<T>& number) {
return to_string(number());
}

template <typename T>
std::ostream& operator<<(std::ostream& stream, const Modular<T>& number) {
return stream << number();
}

template <typename T>
std::istream& operator>>(std::istream& stream, Modular<T>& number) {
typename common_type<typename Modular<T>::Type, int64_t>::type x;
stream >> x;
number.value = Modular<T>::normalize(x);
return stream;
}

/*
using ModType = int;

struct VarMod { static ModType value; };
ModType VarMod::value;
ModType& md = VarMod::value;
using Mint = Modular<VarMod>;

*/

constexpr int md = 998244353;
using Mint = Modular<std::integral_constant<decay<decltype(md)>::type, md>>;

Mint pwr(Mint a, ll k) {
    Mint ans = 1;
    while(k) {
        if(k & 1) ans = ans * a;
        a = a * a;
        k >>= 1;
    }
    return ans;
}

// -- SUFFIX ARRAY START

/*
    Characters should be in range (0, char_bound - 1)
    For strings, use char_bound = 256
*/

template <typename T>
vector<int> suffix_array(int n, const T &s, int char_bound) {
vector<int> a(n);
if (n == 0) {
    return a;
}
if (char_bound != -1) {
    vector<int> aux(char_bound, 0);
    for (int i = 0; i < n; i++) {
    aux[s[i]]++;
    }
    int sum = 0;
    for (int i = 0; i < char_bound; i++) {
    int add = aux[i];
    aux[i] = sum;
    sum += add;
    }
    for (int i = 0; i < n; i++) {
    a[aux[s[i]]++] = i;
    }
} else {
    iota(a.begin(), a.end(), 0);
    sort(a.begin(), a.end(), [&s](int i, int j) { return s[i] < s[j]; });
}
vector<int> sorted_by_second(n);
vector<int> ptr_group(n);
vector<int> new_group(n);
vector<int> group(n);
group[a[0]] = 0;
for (int i = 1; i < n; i++) {
    group[a[i]] = group[a[i - 1]] + (!(s[a[i]] == s[a[i - 1]]));
}
int cnt = group[a[n - 1]] + 1;
int step = 1;
while (cnt < n) {
    int at = 0;
    for (int i = n - step; i < n; i++) {
    sorted_by_second[at++] = i;
    }
    for (int i = 0; i < n; i++) {
    if (a[i] - step >= 0) {
        sorted_by_second[at++] = a[i] - step;
    }
    }
    for (int i = n - 1; i >= 0; i--) {
    ptr_group[group[a[i]]] = i;
    }
    for (int i = 0; i < n; i++) {
    int x = sorted_by_second[i];
    a[ptr_group[group[x]]++] = x;
    }
    new_group[a[0]] = 0;
    for (int i = 1; i < n; i++) {
    if (group[a[i]] != group[a[i - 1]]) {
        new_group[a[i]] = new_group[a[i - 1]] + 1;
    } else {
        int pre = (a[i - 1] + step >= n ? -1 : group[a[i - 1] + step]);
        int cur = (a[i] + step >= n ? -1 : group[a[i] + step]);
        new_group[a[i]] = new_group[a[i - 1]] + (pre != cur);
    }
    }
    swap(group, new_group);
    cnt = group[a[n - 1]] + 1;
    step <<= 1;
}
return a;
}
template <typename T>
vector<int> suffix_array(const T &s, int char_bound) {
return suffix_array((int) s.size(), s, char_bound);
}
template <typename T>
vector<int> build_lcp(int n, const T &s, const vector<int> &sa) {
assert((int) sa.size() == n);
vector<int> pos(n);
for (int i = 0; i < n; i++) {
    pos[sa[i]] = i;
}
vector<int> lcp(max(n - 1, 0));
int k = 0;
for (int i = 0; i < n; i++) {
    k = max(k - 1, 0);
    if (pos[i] == n - 1) {
    k = 0;
    } else {
    int j = sa[pos[i] + 1];
    while (i + k < n && j + k < n && s[i + k] == s[j + k]) {
        k++;
    }
    lcp[pos[i]] = k;
    }
}
return lcp;
}
template <typename T>
vector<int> build_lcp(const T &s, const vector<int> &sa) {
return build_lcp((int) s.size(), s, sa);
}

template<typename T, size_t N>
struct SuffixArray {
    vector<int> sa, lcp;
    int n;
    void init(int nn, const T &s, int char_bound) {
        assert(N >= sz(s));
        n = nn;
        sa = suffix_array(s, char_bound);
        lcp = build_lcp(s, sa);
    }
};

// -- SUFFIX ARRAY END

// -- SPARSE SEGTREE START 

const int NN = 4e7 + 5;
struct node {
    int l, r;
    pair<Mint, Mint> val;
} tr[NN];
int tot = 0;

const int N = 5e5 + 5;
int version[N];
SuffixArray<string, N> S;

Mint pre[N], w[N];
int n;

pair<Mint, Mint> query(int nd,int s,int e,int l,int r) { // sum query in segtree
    if(s > r || e < l || !nd) return {0, 0};
    else if(l <= s && e <= r) {
        return tr[nd].val;
    }
    int m = (s + e) >> 1;
    auto p1 = query(tr[nd].l, s, m, l, r);
    auto p2 = query(tr[nd].r, m + 1, e, l, r);
    return {p1.fi + p2.fi, p1.se + p2.se};
}

void merge(int &nd1,int nd2) { // merge two segtrees
    if(!nd1 || !nd2) {
        nd1 |= nd2;
        return;
    }
    tr[nd1].val.fi += tr[nd2].val.fi;
    tr[nd1].val.se += tr[nd2].val.se;
    merge(tr[nd1].l, tr[nd2].l);
    merge(tr[nd1].r, tr[nd2].r);
}

void upd(int &nd,int s,int e,int idx) { // for initializing segtrees
    if(!nd) nd = ++tot;
    tr[nd].val = {1, pre[idx - 1]};
    if(s == e) return;
    int m = (s + e) >> 1;
    if(idx <= m)
        upd(tr[nd].l, s, m, idx);
    else
        upd(tr[nd].r, m + 1, e, idx);
}

// -- SPARSE SEGTREE END

// -- DSU START

int dsu[N];
Mint cur = 0;
vector<int> pos[N]; // maintains the indexes of the suffix array in each component

int root(int u) {
    return ((u == dsu[u]) ? u : dsu[u] = root(dsu[u]));
}

void uni(int u,int v) { // for merging two components
    u = root(u), v = root(v);
    if(sz(pos[u]) > sz(pos[v])) swap(u, v);
    dsu[u] = v;
    for(int i : pos[u]) {
        auto p1 = query(version[v], 1, n, i + 1, n);
        cur += p1.se - p1.fi * pre[i - 1];
        auto p2 = tr[version[v]].val;
        p2.fi -= p1.fi;
        p2.se -= p1.se;
        cur += p2.fi * pre[i - 1] - p2.se;
        pos[v].pb(i);
    }
    pos[u].clear();
    merge(version[v], version[u]); // merge the second segtree into the first
}

// -- DSU END

void solve() {
    cur = 0;
    while(tot > 0) {
        tr[tot].l = tr[tot].r = 0;
        tr[tot].val = {0, 0};
        tot--;
    }
    string s;
    sc(s);
    n = sz(s);
    S.init(sz(s), s, 256);
    fr(i, 0, 25) sc(w[i]);
    fr(i, 1, n) pre[i] = pre[i - 1] + w[s[i - 1] - 'a'];
    fr(i, 1, n) {
        dsu[i] = i;
        pos[i] = {S.sa[i - 1] + 1};
        version[i] = 0;
        upd(version[i], 1, n, S.sa[i - 1] + 1);
    }
    vector<pair<int,int>> v;
    fr(i, 0, sz(S.lcp) - 1) v.eb(S.lcp[i], i + 1);
    sort(all(v));
    reverse(all(v));
    int pt = 0;
    Mint ans = 0;
    rf(l, n, 1) {
        while(pt < sz(v) && v[pt].fi == l) {
            uni(v[pt].se, v[pt].se + 1);
            pt++;
        }
        ans += cur;
    }
    Mint sum = 0;
    fr(i, 1, n) { // for handling |R| = |T|
        ans += pre[i] * i - sum;
        sum += pre[i];
    }
    ans = ((ans * 2) / (n * 1LL * (n + 1)));
    cout << ans << endl;
}

int main() {
    ios :: sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    int t = 1;
    cin >> t;
    for(int tt = 1; tt <= t; tt++) {
        solve();
    }
    return 0;
}
6 Likes

Why is the time complexity of merging segment trees in the editorial is O(nlog^2n) , isn’t O(nlogn) ?

2 Likes

Merging is O(Nlogn). The extra log accounts for the DSU.

But it only merges n-1 times, so why not O(n(logn+logn)).

1 Like

Yeah, so when we are merging the two segtrees, the complexity of the merge is equal to the number of nodes in the smaller segtree. The number of nodes in the smaller segtree is equal to the log(N) times the number of updates. And merging two segtrees is almost the same as heavy-light merging. So an extra log(N) factor comes due to that. One log factor due to heavy-light merging, one due to the segtree structure. Hope that makes it clear!

I found that your query operation is independent of the merge segment tree. In my code, the answer is updated during the merging of segment tree. So this problem can be solved by O(nlogn) . This is my solution:

https://www.codechef.com/viewsolution/35546365

1 Like

@aviroop123 @enchom Poor Testcases For subtask2 where N is of the order of N<=2000. for N=2000 string=[a…2000times] (i.e. for(2000 times){string+=‘a’}). total number of operations will be of the order of 8*10^9 ~ 10^10. https://www.codechef.com/viewsolution/35606238… this is my solution which got accepted after contest. There are also many partipants whom got AC for 30 points had solution of the order of O(N^3) . Whole time i’m looking for a solution of O(N^2) for 30 points. After Looking the solution of many participants whom got AC I’m disappointed with the Setter & Tester of this problem. If you specify N<=2000 for subtastk2 the solution must be of the order of O(N^2) or O(N^2*logN). Don’t make Poor testcases for a Problem.

3 Likes

What I’m saying is, even without the query function, the complexity would still be O(nlog^2n). I understand that the running time will be much faster, but the complexity will still be the same.

I’m really sorry that O(N^3) solutions have passed the 2nd subtask. I’ll make sure this doesn’t happens again.
But as far as your claim is concerned regarding the test cases, I can say that the number of operations performed in the code that you attached is > 10^9. As you can see in this submission : https://www.codechef.com/viewsolution/35618715 (your code with some added pragma make it faster), the assert fails that the number of operations is <= 10^9.
But again, I’m really sorry that the testing part was not upto the mark and such cache friendly codes passed the test cases.

1 Like

I haven’t heard about sparse segment trees,Can u please add a link where I can read about It?
I tried to find out on internet but I am only getting articles about segment tree and sparse table.
Or should I search with an alternate name(if any)??

you can read here

2 Likes

Can this be solved using Z - Function?
I was able to score 30 points with Z - function O(n^2). Is it possible to optimise it for larger constraints?

2 Likes

how did you manage to get 30 using z function plz elaborate it bro?

Yeah , even I scored 30 points with Z-function and was wondering the same whether it could be extended for 100 points

1 Like

Basically, when you are calculating z-values for a string you end up calulating the z-values for its prefixes too. And it is easy to prove that when i+z[i]=len ,where len is
where 1<=i<n length of string then prefix of length i is a period of the string.
but while iterating a string length of a prefix is atleast i+1 hence i+z[i]>=i+1
implies whenever you have z[i]>=1 then prefix of length i occurs as a period for all substrings for given string z[i] times in total.
To avoid double counting and consider all the sub strings we do the above operations for all the suffixes of the given string.

Here is my code :Z-algo sol (30 points)

3 Likes

Yep. That’s exactly what I did too!

1 Like