cp-library-cpp

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub suisen-cp/cp-library-cpp

:heavy_check_mark: test/src/string/substring_set/arc097_a.test.cpp

Depends on

Code

#define PROBLEM "https://atcoder.jp/contests/arc097/tasks/arc097_c"

#include <iostream>

#include "library/string/substring_set.hpp"
using suisen::SubstringSet;

int main() {
    std::string s;
    int k;
    std::cin >> s >> k;
    std::cout << SubstringSet(s)[k] << std::endl;
    return 0;
}
#line 1 "test/src/string/substring_set/arc097_a.test.cpp"
#define PROBLEM "https://atcoder.jp/contests/arc097/tasks/arc097_c"

#include <iostream>

#line 1 "library/string/substring_set.hpp"



#line 1 "library/string/suffix_automaton.hpp"



#include <algorithm>
#include <cassert>
#include <deque>
#include <limits>
#include <map>
#include <string>
#include <tuple>
#include <vector>

namespace suisen {

/**
 * reference
 * - https://w.atwiki.jp/uwicoder/pages/2842.html
 * - https://cp-algorithms.com/string/suffix-automaton.html
 */
template <typename T, typename SequenceType>
struct SuffixAutomatonBase {
    struct Node {
        std::map<T, int> nxt;
        int link, len;
        bool cloned;
        int index;
    };

    std::vector<Node> nodes;
    int last;

    SuffixAutomatonBase() {
        nodes.push_back({ {}, -1, 0, false, -1 });
        last = 0;
    }
    SuffixAutomatonBase(const SequenceType &s) : SuffixAutomatonBase() {
        for (const T &c : s) append(c);
    }

    void append(const T &c) {
        const int new_node = nodes.size();
        const int new_node_index = nodes[last].index + 1;
        nodes.push_back({ {}, -1, nodes[last].len + 1, false, new_node_index });
        int p = last;
        for (; p != -1 and not nodes[p].nxt.count(c); p = nodes[p].link) nodes[p].nxt[c] = new_node;
        const int q = p == -1 ? 0 : nodes[p].nxt[c];
        if (p == -1 or nodes[p].len + 1 == nodes[q].len) {
            nodes[new_node].link = q;
        } else {
            const int clone_node = nodes.size();
            nodes.push_back({ nodes[q].nxt, nodes[q].link, nodes[p].len + 1, true, new_node_index });
            for (; p != -1 and nodes[p].nxt[c] == q; p = nodes[p].link) nodes[p].nxt[c] = clone_node;
            nodes[new_node].link = nodes[q].link = clone_node;
        }
        last = new_node;
    }
    SuffixAutomatonBase& operator+=(const T &c) {
        append(c);
        return *this;
    }

    int transition(const SequenceType &t, int invalid_state = -1) const {
        int cur = 0;
        for (const auto &c : t) {
            auto it = nodes[cur].nxt.find(c);
            if (it == nodes[cur].nxt.end()) return invalid_state;
            cur = it->second;
        }
        return cur;
    }

    bool accept(const SequenceType &t) const {
        return transition(t) != -1;
    }

    class SubstringCounter {
        public:
            SubstringCounter(const SuffixAutomatonBase *sa) : sa(sa), n(sa->nodes.size()), dp(n, 0) {
                const std::vector<Node> &nodes = sa->nodes;
                for (const int u : sa->topological_order(/* reversed = */ true)) {
                    dp[u] += not nodes[u].cloned;
                    const int p = nodes[u].link;
                    if (p >= 0) dp[p] += dp[u];
                }
            }
            long long count(const SequenceType &t) const {
                const int state = sa->transition(t);
                return state == -1 ? 0 : dp[state];
            }
        private:
            const SuffixAutomatonBase *sa;
            int n;
            std::vector<long long> dp;
    };

    SubstringCounter substring_counter() const & {
        return SubstringCounter { this };
    }
    SubstringCounter substring_counter() const && = delete;

    class SuffixLinkTree {
        public:
            SuffixLinkTree(const SuffixAutomatonBase *sa) : sa(sa), g(sa->nodes.size()) {
                const int n = g.size();
                for (int i = 1; i < n; ++i) g[sa->nodes[i].link].push_back(i);
            }
            int size() const {
                return g.size();
            }
            const std::vector<int>& operator[](int i) const {
                return g[i];
            }
        private:
            const SuffixAutomatonBase *sa;
            std::vector<std::vector<int>> g;
    };

    SuffixLinkTree suffix_link_tree() const & {
        return SuffixLinkTree(this);
    }
    SuffixLinkTree suffix_link_tree() const && = delete;

    class OccurrenceEnumerator {
        public:
            OccurrenceEnumerator(const SuffixAutomatonBase *sa) : sa(sa), t(sa->suffix_link_tree()) {}

            // returns vector of i s.t. S[i:i+|t|] = t
            std::vector<int> enumerate_all_occurrence(const SequenceType &pattern) const {
                const int state = sa->transition(pattern);
                if (state == -1) return {};
                const std::vector<Node> &nodes = sa->nodes;
                const int l = pattern.size();
                std::vector<int> res;
                auto dfs = [&](auto self, int u) -> void {
                    if (not nodes[u].cloned) res.push_back(nodes[u].len - l);
                    for (const int v : t[u]) self(self, v);
                };
                dfs(dfs, state);
                return res;
            }

        private:
            const SuffixAutomatonBase *sa;
            SuffixLinkTree t;
    };

    OccurrenceEnumerator occurrence_enumerator() const & {
        return OccurrenceEnumerator(this);
    }
    OccurrenceEnumerator occurrence_enumerator() const && = delete;

    class FirstOccurenceSearcher {
        public:
            FirstOccurenceSearcher(const SuffixAutomatonBase *sa) : sa(sa) {
                const std::vector<Node> &nodes = sa->nodes;
                dp.resize(nodes.size(), std::numeric_limits<int>::max());
                for (const int u : sa->topological_order(/* reversed = */ true)) {
                    if (not nodes[u].cloned) dp[u] = nodes[u].len;
                    const int p = nodes[u].link;
                    if (p >= 0 and nodes[p].cloned) dp[p] = std::min(dp[p], dp[u]);
                }
            }

            // returns min { i | S[i:i+|t|] = t }. if such i does not exist, returns -1.
            int first_occurrence(const SequenceType &t) const {
                const int state = sa->transition(t);
                if (state == -1) return -1;
                return dp[state] - t.size();
            }

        private:
            const SuffixAutomatonBase *sa;
            std::vector<int> dp;
    };

    FirstOccurenceSearcher first_occurence_searcher() const & {
        return FirstOccurenceSearcher(this);
    }
    FirstOccurenceSearcher first_occurence_searcher() const && = delete;

    // returns { start_s, start_t, len } s.t. lcs = s[start_s: start_s+len] t[start_t: start_t+len]
    std::tuple<int, int, int> longest_common_substring(const SequenceType &t) const {
        if (t.size() == 0) return { 0, 0, 0 };
        const Node *v = &nodes[0];
        int l = 0, max_len = 0, s_end_pos = 0, t_end_pos = 0;
        for (int i = 0; i < int(t.size()); ++i){
            while (v->link != -1 and not v->nxt.count(t[i])) {
                v = &nodes[v->link];
                l = v->len;
            }
            auto it = v->nxt.find(t[i]);
            if (it != v->nxt.end()){
                v = &nodes[it->second];
                l++;
            }
            if (l > max_len){
                max_len = l;
                t_end_pos = i;
                s_end_pos = v->index;
            }
        }
        if (max_len == 0) return { 0, 0, 0 };
        return { s_end_pos - max_len + 1, t_end_pos - max_len + 1, max_len };
    }

    std::vector<int> topological_order(bool reversed = false) const {
        const int n = nodes.size();
        std::vector<int> in(n, 0);
        for (const auto &node : nodes) {
            for (const auto &p : node.nxt) ++in[p.second];
        }
        std::deque<int> dq;
        for (int i = 0; i < n; ++i) {
            if (in[i] == 0) dq.push_back(i);
        }
        std::vector<int> res;
        while (dq.size()) {
            int u = dq.front();
            dq.pop_front();
            res.push_back(u);
            for (const auto &p : nodes[u].nxt) {
                if (--in[p.second] == 0) dq.push_back(p.second);
            }
        }
        if (reversed) std::reverse(res.begin(), res.end());
        assert(int(res.size()) == n);
        return res;
    }
};

template <typename T>
struct SuffixAutomaton : public SuffixAutomatonBase<T, std::vector<T>> {
    using SuffixAutomatonBase<T, std::vector<T>>::SuffixAutomatonBase;
    using value_type = T;
    using sequence_type = std::vector<T>;
};

template <typename T>
SuffixAutomaton(std::vector<T>) -> SuffixAutomaton<T>;

template <>
struct SuffixAutomaton<char> : public SuffixAutomatonBase<char, std::string> {
    using SuffixAutomatonBase<char, std::string>::SuffixAutomatonBase;
    using value_type = char;
    using sequence_type = std::string;
};

SuffixAutomaton(std::string) -> SuffixAutomaton<char>;

} // namespace suisen


#line 5 "library/string/substring_set.hpp"

namespace suisen {

/**
 * Constant set of all substrings
 */
template <typename T>
class SubstringSet {
    public:
        using value_type = typename SuffixAutomaton<T>::sequence_type;
        using size_type = long long;
        using difference_type = size_type;

        // create empty set
        SubstringSet() : SubstringSet(value_type{}) {}

        // create set of all substrings in `s`
        SubstringSet(const value_type &s) {
            build(s);
        }

        // build set of all substrings in `s`
        void build(const value_type &s) {
            sa = SuffixAutomaton(s);
            dp.resize(sa.nodes.size(), size_type(1));
            for (const int u : sa.topological_order(/* reversed = */true)) {
                for (const auto &p : sa.nodes[u].nxt) dp[u] += dp[p.second];
            }
        }

        size_type size() const {
            return dp[0];
        }

        bool contains(const value_type &t) const {
            return sa.accept(t);
        }

        // the k'th lexicographically smallest substring, O(|return value|).
        value_type operator[](size_type k) const {
            assert(0 <= k and k < dp[0]);
            int cur = 0;
            value_type res;
            while (k--) {
                for (const auto &[e, v] : sa.nodes[cur].nxt) {
                    if (k < dp[v]) {
                        res.push_back(e);
                        cur = v;
                        break;
                    } else {
                        k -= dp[v];
                    }
                }
            }
            return res;
        }
        // the k'th lexicographically smallest substring, O(|return value|).
        value_type kth_element(size_type k) const {
            return (*this)[k];
        }

        // #{ s | s < t, s is in this set }, O(|min(t, max_substr_len)|).
        size_type count_lt(const value_type &t) const {
            size_type res = 0;
            int cur = 0;
            for (const T& c : t) {
                ++res;
                auto it_r = sa.nodes[cur].nxt.lower_bound(c);
                for (auto it_l = sa.nodes[cur].nxt.begin(); it_l != it_r; ++it_l) {
                    res += dp[it_l->second];
                }
                if (it_r == sa.nodes[cur].nxt.end() or it_r->first != c) break;
                cur = it_r->second;
            }
            return res;
        }
        // #{ s | s <= t, s is in this set }, O(min(|t|, max_substr_len)).
        size_type count_leq(const value_type &t) const {
            return count_lt(t) + contains(t);
        }
        // #{ s | s > t, s is in this set }, O(min(|t|, max_substr_len)).
        size_type count_gt(const value_type &t) const {
            return size() - count_leq(t);
        }
        // #{ s | s >= t, s is in this set }, O(min(|t|, max_substr_len)).
        size_type count_geq(const value_type &t) const {
            return size() - count_lt(t);
        }
        // #{ s | l <= s <= r, s is in this set }, O(min(|t|, max_substr_len)).
        size_type count_range(const value_type &l, const value_type &r) const {
            return l >= r ? 0 : count_lt(r) - count_lt(l);
        }

        // min { s | s >= t, s is in this set }, O(|return value| + min(|t|, max_substr_len))
        value_type min_geq(const value_type &t) const {
            return (*this)[count_lt(t)];
        }
        // min { s | s > t, s is in this set }, O(|return value| + min(|t|, max_substr_len))
        value_type min_gt(const value_type &t) const {
            return (*this)[count_leq(t)];
        }
        // max { s | s < t, s is in this set }, O(|return value| + min(|t|, max_substr_len))
        value_type max_lt(const value_type &t) const {
            return (*this)[count_lt(t) - 1];
        }
        // max { s | s <= t, s is in this set }, O(|return value| + min(|t|, max_substr_len))
        value_type max_leq(const value_type &t) const {
            return (*this)[count_leq(t) - 1];
        }

        // iterator

        // operator*: O(|return value|), other operations: O(1).
        class SubstringSetIterator {
            public:
                SubstringSetIterator(const SubstringSet *st, size_type k) : st(st), k(k) {}

                value_type operator*() const { return st->kth_substring(k); }

                auto& operator=(const SubstringSetIterator& other) {
                    st = other.st;
                    k = other.k;
                    return *this;
                }

                auto& operator++() { ++k; return *this; }
                auto& operator--() { --k; return *this; }
                auto& operator+=(difference_type d) { k += d; return *this; }
                auto& operator-=(difference_type d) { k -= d; return *this; }

                auto operator+(difference_type d) { return SubstringSetIterator { st, k + d }; }
                auto operator-(difference_type d) { return SubstringSetIterator { st, k - d }; }

                difference_type operator-(const SubstringSetIterator &other) { return k - other.k; }

                bool operator==(const SubstringSetIterator& other) const { return k == other.k; }
                bool operator!=(const SubstringSetIterator& other) const { return k != other.k; }
                bool operator< (const SubstringSetIterator& other) const { return k <  other.k; }
                bool operator<=(const SubstringSetIterator& other) const { return k <= other.k; }
                bool operator> (const SubstringSetIterator& other) const { return k >  other.k; }
                bool operator>=(const SubstringSetIterator& other) const { return k >= other.k; }

            protected:
                const SubstringSet *st;
                size_type k;
        };

        // operator*: O(|return value|), other operations: O(1).
        struct ReversedSubstringSetIterator : public SubstringSetIterator {
            using SubstringSetIterator::SubstringSetIterator;
            value_type operator*() const { return this->st->kth_element(this->st->size() - this->k - 1); }
        };

        using iterator = SubstringSetIterator;
        using reverse_iterator = ReversedSubstringSetIterator;
        using const_iterator = iterator;
        using const_reverse_iterator = reverse_iterator;

        auto begin()   const { return iterator { this, 0LL }; }
        auto end()     const { return iterator { this, size() }; }
        auto cbegin()  const { return begin(); }
        auto cend()    const { return end(); }
        auto rbegin()  const { return reverse_iterator { this, 0LL }; }
        auto rend()    const { return reverse_iterator { this, size() }; }
        auto crbegin() const { return rbegin(); }
        auto crend()   const { return rend(); }

        // O(|min(t, max_substr_len)|).
        auto lower_bound(const value_type &t) const {
            return iterator { this, count_lt(t) };
        }
        // O(|min(t, max_substr_len)|).
        auto upper_bound(const value_type &t) const {
            return iterator { this, count_leq(t) };
        }

        // O(|min(t, max_substr_len)|).
        auto find(const value_type &t) const {
            auto it = lower_bound(t);
            if (it == end() or t != *it) return end();
            return it;
        }

    private:
        SuffixAutomaton<T> sa;
        std::vector<size_type> dp;
};

template <typename T>
SubstringSet(std::vector<T>) -> SubstringSet<T>;

SubstringSet(std::string) -> SubstringSet<char>;

} // namespace suisen



#line 6 "test/src/string/substring_set/arc097_a.test.cpp"
using suisen::SubstringSet;

int main() {
    std::string s;
    int k;
    std::cin >> s >> k;
    std::cout << SubstringSet(s)[k] << std::endl;
    return 0;
}
Back to top page