はまやんはまやんはまやん

hamayanhamayan's blog

SubstringQueries [Single Round Match 779 Round 1 - Division I, Level Two Med]

競技プログラミング

競プロのコード例について補足 競技プログラミング練習問題集 Twitter 競技セキュリティまとめのまとめ

https://community.topcoder.com/stat?c=problem_statement&pm=15740

前提知識

解説

テクニックとして、ある文字列の任意の2区間の大小関係は前計算しておけば、O(1)で判定することができる。
Suffix Array/LCPでそれが実現できる。
これを持っていないと以下の方針で解くのは難しい。

簡単な別解を補足で載せますが、ここは本番回答。

Nを文字列長としておく。

Kで場合分けして解く。
K=1の時はf(X)を計算可能なので計算しよう。
全部でO(N2)で計算できる。

2≦Kの場合は、性質を見抜く必要がある。
(1) 1≦X≦Nの場合は、Xを大きくしていくと、Q(X)も広義単調増加していく。
(2) そこからXを大きくしていくと、Q(X)の値は変わらず、区間が伸びていく。
(3) 区間が文字列の右端に達したとき、Q(X)が広義単調減少していく。
(1)はN通りしかないし、(2)は区間を伸ばせるだけ伸ばすので1通りだけ考えればよく、
(3)はQ(X)の変化点だけを計算すると高々N通りしかない。
どのパートにおいても、O(N2)で計算可能なので、これをまとめると答え。

#ifdef _MSC_VER
inline unsigned int __builtin_clz(unsigned int x) { unsigned long r; _BitScanReverse(&r, x); return 31 - r; }
#endif // _MSC_VER
template<class V> struct SparseTable { // [L,R)
    const V def = 0;
    inline V comp(V a, V b) { return min(a, b); }

    int n; vector<V> a, b[22]; inline int __lg(int x) { return 32 - 1 - __builtin_clz(x); }
    void init(vector<V> v) {
        int nn = v.size(); n = 1; while (n < nn) n *= 2; a.resize(n);
        rep(i, 0, 22) b[i].resize(n); rep(i, 0, nn) a[i] = v[i];

        int d = 1 << __lg(n - 1), e = d << 1;
        for (int h = 0, f; (f = 1 << h) <= d; ++h) {
            for (int i = f, r = f << 1; i < e; i += r) {
                b[h][i - 1] = a[i - 1];
                for (int j = i - 2; j >= i - f; --j) b[h][j] = comp(b[h][j + 1], a[j]);
                b[h][i] = a[i];
                for (int j = i + 1; j < i + f; ++j) b[h][j] = comp(b[h][j - 1], a[j]);
            }
        }
    }

    V get(int L, int R) {
        assert(0 <= L && L <= R); if (L == R)return def; R--; if (L == R)return a[L]; int h = __lg(L ^ R);
        return comp(b[h][L], b[h][R]);
    }
};
struct SuffixArraySAIS {
    vector<int> sa, lcp, rank;
    SparseTable<int> rmq;
    SuffixArraySAIS() {}
    SuffixArraySAIS(string str_) : str(str_) {
        init(str_);
    }
    SuffixArraySAIS(const vector<int>& S_, int A_SIZE_, bool lcp_flag = true) : S(S_), A_SIZE(A_SIZE_) {
        buildSA();
        if (lcp_flag) {
            buildLCP();
            buildRMQ();
        }
    }
    void init(string str_) {
        str = str_;
        N = str.size() + 1;
        S = vector<int>(N, 0);
        for (int i = 0; i < N; i++) S[i] = str[i];
        A_SIZE = 256;

        buildSA();
        buildLCP();
        buildRMQ();
    }
    string str;
    vector<int> S;
    int A_SIZE;
    int N;
    vector<int> t, B;
    enum { STYPE, LTYPE };

    inline bool is_lms(int i) {
        return i > 0 && t[i] == STYPE && t[i - 1] == LTYPE;
    }
    void bucket() {
        B = vector<int>(A_SIZE);
        for (int i = 0; i < N; i++) B[S[i]]++;
        for (int i = 0; i < A_SIZE - 1; i++) B[i + 1] += B[i];
    }
    void induced_sort() {
        bucket();
        for (int i = 0; i < N; i++) {
            int j = sa[i] - 1;
            if (j >= 0 && S[j] >= S[j + 1]) sa[B[S[j] - 1]++] = j;
        }
        bucket();
        for (int i = N; i--; ) {
            int j = sa[i] - 1;
            if (j >= 0 && S[j] <= S[j + 1]) sa[--B[S[j]]] = j;
        }
    }
    void buildSA() {
        N = S.size();
        sa.assign(N, -1);
        if (N == 1) {
            sa[0] = 0;
            return;
        }
        t.assign(N, STYPE);
        for (int i = N - 1; i--;)
            if (S[i] > S[i + 1] || (S[i] == S[i + 1] && t[i + 1] == LTYPE))
                t[i] = LTYPE;
        bucket();
        for (int i = N; i--;)
            if (is_lms(i)) sa[--B[S[i]]] = i;
        induced_sort();

        int N1 = 0;
        for (int i = 0; i < N; i++) if (is_lms(sa[i])) sa[N1++] = sa[i];

        fill(sa.begin() + N1, sa.end(), -1);
        int name = 0, prev = -1;
        for (int i = 0; i < N1; i++) {
            int cur = sa[i];
            bool diff = (prev == -1);
            for (int j = 0; !diff; j++) {
                if (j > 0 && is_lms(cur + j)) break;
                if (S[cur + j] != S[prev + j]) diff = true;
            }
            if (diff) name++;
            sa[N1 + cur / 2] = name - 1;
            prev = cur;
        }
        vector<int> S1, sa1(N1);
        for (int i = N1; i < N; i++) if (sa[i] >= 0) S1.push_back(sa[i]);
        if (name == N1) for (int i = 0; i < N1; i++) sa1[S1[i]] = i;
        else sa1 = SuffixArraySAIS(S1, name, false).sa;

        N1 = 0;
        for (int i = 0; i < N; i++) if (is_lms(i)) S1[N1++] = i;
        for (int i = 0; i < N1; i++) sa1[i] = S1[sa1[i]];

        fill(sa.begin(), sa.end(), -1);
        bucket();
        for (int i = N1; i--;) {
            int j = sa1[i];
            sa[--B[S[j]]] = j;
        }
        induced_sort();
    }
    void buildLCP() {
        rank.resize(N);
        lcp.resize(N - 1);
        for (int i = 0; i < N; i++) rank[sa[i]] = i;
        int h = 0;
        for (int i = 0; i < N - 1; i++) {
            int j = sa[rank[i] - 1];
            if (h > 0) h--;
            for (; j + h < N && i + h < N && S[j + h] == S[i + h]; h++);
            lcp[rank[i] - 1] = h;
        }
    }

    void buildRMQ() {
        rmq.init(lcp);
    }
    int common_prefix(int x, int y) {
        if (x == y) return N - 1 - x;
        if (y >= N - 1) return 0;
        if (rank[x] > rank[y]) swap(x, y);
        return rmq.get(rank[x], rank[y]);
    }
    int compare(int x, int xn, int y, int yn) {
        int l = common_prefix(x, y);
        if (l >= xn || l >= yn) return xn < yn ? -1 : xn == yn ? 0 : 1;
        return rank[x] < rank[y] ? -1 : x == y ? 0 : 1;
    }
};

ll mo = 1;

struct SubstringQueries {
    long long solve(string S, int k) {
        SuffixArraySAIS sa(S + S);
        int n = S.length();

        rep(i, 0, 18) mo *= 10;

        if (k == 1) {
            ll res = 0;
            rep(len, 1, n + 1) {
                int top = 0;
                rep(i, 0, n - len + 1) {
                    if (sa.compare(top, len, i, len) > 0) {
                        top = i;
                    }
                }
                res = (res + top) % mo;
            }
            return res;
        }

        ll res = 0;
        int pre = 0;
        rep(len, 1, n + 1) {
            int top = pre;
            rep(i, pre, n) {
                if (sa.compare(top, len, i, len) > 0) {
                    top = i;
                }
            }
            res = (res + top) % mo;
            pre = top;
        }

        ll len = n;
        ll malen = 1LL * n * k;

        res = (res + 1LL * pre * (malen - len - pre)) % mo;
        len = pre;

        while (0 < pre) {
            // [0, pre)
            int mi = 0;
            rep(i, 0, pre) if (sa.compare(mi, n, i, n) >= 0) mi = i;
            res = (res + 1LL * mi * (pre - mi)) % mo;
            pre = mi;
        }

        return res;
    }
};

簡単な別解

snukeさんのツイートにk=2,3の答えを求めたら、あとは等差数列とあり、なるほどとなった解法。
上の解法にもあるが、kが増えていくと、(2) そこからXを大きくしていくと、Q(X)の値は変わらず、区間が伸びていく。 の部分だけ伸びていく。
Q(X)はkの値によらず、kを1つ増やすと、この区間がNだけ増える。よって同じ分だけ増えてくので等差数列として考えられる。
なるほどすぎるな。
実装もSuffixArray/LCPがすでになされていればかなり簡潔。

struct SubstringQueries {
    ll mo = 1;
    SuffixArraySAIS sa;

    ll naive(string S) {
        int n = S.length();
        sa.init(S);

        ll res = 0;
        rep(x, 1, n + 1) {
            int qx = 0;
            rep(i, 0, n - x + 1) {
                if (sa.compare(qx, x, i, x) > 0) {
                    qx = i;
                }
            }
            res = (res + qx) % mo;
        }
        return res;
    }

    long long solve(string S, int k) {
        rep(i, 0, 18) mo *= 10;

        if (k <= 3) {
            string s = "";
            rep(i, 0, k) s += S;
            return naive(s);
        }

        ll k2 = naive(S + S);
        ll k3 = naive(S + S + S);
        ll d = k3 - k2;

        return (k2 + d * (k - 2)) % mo;
    }
};