string/sa_lcp.hpp

• View this file on GitHub
• Last update: 2024-07-18 16:56:22+09:00
• Include: #include "string/sa_lcp.hpp"

Depends on

• data_structure/sparse_table.hpp

Code

#pragma once
#include <atcoder/string>
#include "../data_structure/sparse_table.hpp"

template <typename Cont>
struct StringSaLcp {
    Cont s;
    std::vector<int> sa, lcp, rev;
    SparseTable<int> sp;

    StringSaLcp(const Cont &s)
        : s(s),
          sa(atcoder::suffix_array(s)),
          lcp(atcoder::lcp_array(s, sa)),
          rev(s.size()),
          sp(lcp) {
        for (int i = 0; i < (int)s.size(); ++i) {
            rev[sa[i]] = i;
        }
    }

    // lcp(s[i:], s[j:])
    int get_lcp(int i, int j) const {
        assert(0 <= i && i <= (int)s.size());
        assert(0 <= j && j <= (int)s.size());
        if (i == (int)s.size() || j == (int)s.size()) {
            return 0;
        }
        if (i == j) {
            return (int)s.size() - i;
        }
        i = rev[i];
        j = rev[j];
        if (i > j) {
            std::swap(i, j);
        }
        return sp.min(i, j);
    }

    // lcp(s[l0:r0], s[l1:r1])
    int get_lcp(int l0, int r0, int l1, int r1) const {
        assert(0 <= l0 && l0 <= r0 && r0 <= (int)s.size());
        assert(0 <= l1 && l1 <= r1 && r1 <= (int)s.size());
        if (l0 == r0 || l1 == r1) {
            return 0;
        }
        return std::min({get_lcp(l0, l1), r0 - l0, r1 - l1});
    }

    int cmp(int l0, int r0, int l1, int r1) const {
        assert(0 <= l0 && l0 <= r0 && r0 <= (int)s.size());
        assert(0 <= l1 && l1 <= r1 && r1 <= (int)s.size());
        int lc = get_lcp(l0, r0, l1, r1);
        if (lc == r0 - l0 && lc == r1 - l1) {
            return 0;
        }
        if (lc == r0 - l0 || (lc != r1 - l1 && s[l0 + lc] < s[l1 + lc])) {
            return -1;
        } else {
            return 1;
        }
    }
};

#line 2 "string/sa_lcp.hpp"
#include <atcoder/string>
#line 2 "data_structure/sparse_table.hpp"

#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>
#include <functional>

template <typename T, typename F = std::less<T>>
class SparseTable {
    std::vector<std::vector<T>> table;
    int s;
    const F f;

    int log2(int n) const {
        return 31 - __builtin_clz(n);
    }
    
    T min2(const T &x, const T &y) const {
        if (f(x, y)) {
            return x;
        } else {
            return y;
        }
    }

public:
    SparseTable(std::vector<T> arr, const F &f = F()) : s((int) arr.size()), f(f) {
        if (s == 0) {
            return;
        }
        int m = log2(s);
        table.resize(m + 1);
        table[0] = std::move(arr);
        for (int i = 1; i <= m; ++i) {
            int w = 1 << i;
            table[i].resize(s - w + 1);
            for (int j = 0; j < (int) table[i].size(); ++j) {
                table[i][j] = min2(table[i - 1][j], table[i - 1][j + (w >> 1)]);
            }
        }
    }

    int size() const {
        return s;
    }

    // not empty
    T min(int l, int r) const {
        assert(l >= 0 && l < r && r <= s);
        int m = log2(r - l);
        return min2(table[m][l], table[m][r - (1 << m)]);
    }
};

#line 4 "string/sa_lcp.hpp"

template <typename Cont>
struct StringSaLcp {
    Cont s;
    std::vector<int> sa, lcp, rev;
    SparseTable<int> sp;

    StringSaLcp(const Cont &s)
        : s(s),
          sa(atcoder::suffix_array(s)),
          lcp(atcoder::lcp_array(s, sa)),
          rev(s.size()),
          sp(lcp) {
        for (int i = 0; i < (int)s.size(); ++i) {
            rev[sa[i]] = i;
        }
    }

    // lcp(s[i:], s[j:])
    int get_lcp(int i, int j) const {
        assert(0 <= i && i <= (int)s.size());
        assert(0 <= j && j <= (int)s.size());
        if (i == (int)s.size() || j == (int)s.size()) {
            return 0;
        }
        if (i == j) {
            return (int)s.size() - i;
        }
        i = rev[i];
        j = rev[j];
        if (i > j) {
            std::swap(i, j);
        }
        return sp.min(i, j);
    }

    // lcp(s[l0:r0], s[l1:r1])
    int get_lcp(int l0, int r0, int l1, int r1) const {
        assert(0 <= l0 && l0 <= r0 && r0 <= (int)s.size());
        assert(0 <= l1 && l1 <= r1 && r1 <= (int)s.size());
        if (l0 == r0 || l1 == r1) {
            return 0;
        }
        return std::min({get_lcp(l0, l1), r0 - l0, r1 - l1});
    }

    int cmp(int l0, int r0, int l1, int r1) const {
        assert(0 <= l0 && l0 <= r0 && r0 <= (int)s.size());
        assert(0 <= l1 && l1 <= r1 && r1 <= (int)s.size());
        int lc = get_lcp(l0, r0, l1, r1);
        if (lc == r0 - l0 && lc == r1 - l1) {
            return 0;
        }
        if (lc == r0 - l0 || (lc != r1 - l1 && s[l0 + lc] < s[l1 + lc])) {
            return -1;
        } else {
            return 1;
        }
    }
};

Back to top page