test/aoj/2292.test.cpp

View this file on GitHub
Last update: 2024-01-07 23:25:49+09:00
Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2292

Depends on

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2292"

#include "../../string/palindromic_tree.hpp"
#include "../../string/rolling_hash.hpp"

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout << fixed << setprecision(15);

    string s, t;
    cin >> s >> t;
    PalindromicTree pts(s), ptt(t);
    ll base = RollingHash::generate_base();
    RollingHash rhs(s, base), rht(t, base);
    map<ll, ll> mp;
    for (auto [len, idx, cnt] : pts.get_palindrome_frequencies()) {
        mp[rhs.query(idx, idx+len)] += cnt;
    }
    ll ans = 0;
    for (auto [len, idx, cnt] : ptt.get_palindrome_frequencies()) {
        ans += mp[rht.query(idx, idx+len)] * cnt;
    }
    cout << ans << endl;
}

#line 1 "test/aoj/2292.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2292"

#line 2 "string/palindromic_tree.hpp"
#include <map>
#include <ranges>
#include <string>
#include <vector>

class PalindromicTree {
   public:
    PalindromicTree() : nodes(2) {
        nodes[0] = Node(-1, -1, 0);  // root for odd length palindromes
        nodes[1] = Node(0, -1, 0);   // root for even length palindromes
        nodes[1].suffix_link = 0;
        suff = 0;
    }

    explicit PalindromicTree(const std::string& s) : PalindromicTree() {
        for (char c : s) add(c);
    }

    void add(char c) {
        str.push_back(c);

        //  find a palidrome cAc
        int k = find_next_palindrome(suff);

        // the palindrome already exists
        if (nodes[k].link.contains(c)) {
            ++nodes[nodes[k].link[c]].cnt;
            suff = nodes[k].link[c];
            return;
        }

        // create a new node
        nodes[k].link[c] = nodes.size();
        suff = nodes.size();
        nodes.emplace_back(nodes[k].len + 2, (int)str.size() - nodes[k].len - 2,
                           1);

        // add a suffix link
        if (nodes.back().len == 1) {
            nodes.back().suffix_link = 1;
        } else {
            const int n = find_next_palindrome(nodes[k].suffix_link);
            nodes.back().suffix_link = nodes[n].link[c];
        }
    }

    std::vector<int> get_suffix_palindromes() const {
        std::vector<int> ret;
        int k = suff;
        while (nodes[k].len > 0) {
            ret.push_back(nodes[k].len);
            k = nodes[k].suffix_link;
        }
        return ret;
    }

    // returns {length, one of the starting indices, count}
    std::vector<std::tuple<int, int, int>> get_palindrome_frequencies() {
        std::vector<std::tuple<int, int, int>> ret;
        for (auto& node : nodes | std::views::drop(2) | std::views::reverse) {
            ret.emplace_back(node.len, node.idx, node.cnt);
            nodes[node.suffix_link].cnt += node.cnt;
        }
        return ret;
    }

   private:
    struct Node {
        std::map<char, int> link;
        int suffix_link;
        int len, idx, cnt;

        Node() = default;
        Node(int len, int idx, int cnt) : len(len), idx(idx), cnt(cnt) {}
    };

    std::vector<Node> nodes;
    int suff;
    std::string str;

    int find_next_palindrome(int k) const {
        const int pos = str.size() - 1;
        while (true) {
            int i = pos - 1 - nodes[k].len;
            if (i >= 0 && str[i] == str[pos]) break;
            k = nodes[k].suffix_link;
        }
        return k;
    }
};
#line 2 "string/rolling_hash.hpp"
#include <random>
#line 5 "string/rolling_hash.hpp"

class RollingHash {
   public:
    static long long generate_base() {
        std::random_device rd;
        std::mt19937_64 rng(rd());
        std::uniform_int_distribution<long long> rand(1, mod - 1);
        return rand(rng);
    }

    RollingHash() = default;
    RollingHash(const std::string& s, long long base)
        : RollingHash(std::vector<char>(s.begin(), s.end()), base) {}
    template <typename T>
    RollingHash(const std::vector<T>& s, long long base)
        : base(base), hashed(s.size() + 1), power(s.size() + 1) {
        power[0] = 1;
        for (int i = 0; i < (int)s.size(); ++i) {
            power[i + 1] = mul(power[i], base);
            hashed[i + 1] = add(mul(hashed[i], base), s[i]);
        }
    }

    long long query(int l, int r) const {
        return add(hashed[r], mod - mul(hashed[l], power[r - l]));
    }

    long long combine(long long h1, long long h2, int len2) const {
        return add(mul(h1, power[len2]), h2);
    }

    void push_back(char c) {
        power.push_back(mul(power.back(), base));
        hashed.push_back(add(mul(hashed.back(), base), c));
    }

   private:
    static constexpr long long mod = (1LL << 61) - 1;

    static inline long long add(long long a, long long b) {
        if ((a += b) >= mod) a -= mod;
        return a;
    }

    static inline long long mul(long long a, long long b) {
        __int128_t c = (__int128_t)a * b;
        return add(c >> 61, c & mod);
    }

    const long long base;
    std::vector<long long> hashed, power;
};
#line 5 "test/aoj/2292.test.cpp"

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout << fixed << setprecision(15);

    string s, t;
    cin >> s >> t;
    PalindromicTree pts(s), ptt(t);
    ll base = RollingHash::generate_base();
    RollingHash rhs(s, base), rht(t, base);
    map<ll, ll> mp;
    for (auto [len, idx, cnt] : pts.get_palindrome_frequencies()) {
        mp[rhs.query(idx, idx+len)] += cnt;
    }
    ll ans = 0;
    for (auto [len, idx, cnt] : ptt.get_palindrome_frequencies()) {
        ans += mp[rht.query(idx, idx+len)] * cnt;
    }
    cout << ans << endl;
}