add palindromic tree

Author: hitonanode

string/palindromic_tree.hpp

@@ -0,0 +1,97 @@
+#pragma once
+
+#include <map>
+#include <string>
+#include <vector>
+
+// Palindromic tree / Eertree （回文木）
+namespace palindromic_tree {
+
+template <class Key> class Node {
+    int suffix_link_; // このノードからのsuffix link （suffix の最長回文）
+    int length_; // このノードが表す回文の長さ。 -1 となる場合もあるので注意
+    std::map<Key, int> children;
+
+public:
+    explicit Node(int suffix_link, int length) : suffix_link_(suffix_link), length_(length) {}
+
+    int suffix_link() const { return suffix_link_; }
+
+    int length() const { return length_; }
+
+    int get_child(Key c) const {
+        auto it = children.find(c);
+        return (it == children.end()) ? -1 : it->second;
+    }
+
+    void set_child(int c, int nxt_idx) { children[c] = nxt_idx; }
+
+    template <class OStream> friend OStream &operator<<(OStream &os, const Node &node) {
+        os << "Node(suffix_link=" << node.suffix_link() << ", length=" << node.length()
+           << ", children={";
+        for (const auto &[c, nxt] : node.children) os << c << "->" << nxt << ", ";
+        return os << "})";
+    }
+};
+
+// Palindromic tree
+// nodes[0] は長さ -1, nodes[1] は長さ 1 のダミーノード
+template <class Key> struct Tree {
+    std::vector<Node<Key>> nodes;
+
+    Tree() { nodes = {Node<Key>(-1, -1), Node<Key>(0, 0)}; }
+
+    // nodes[cursor] は s[0:i] の suffix palindrome を表す
+    // 本関数はその nodes[cursor] の suffix palindrome であって更に s[0:(i + 1)] の suffix link となりうる最長のものを返す
+    int find_next_suffix(const std::vector<Key> &s, int i, int cursor) {
+        while (true) {
+            if (cursor < 0) return 0;
+
+            const int cur_len = nodes.at(cursor).length();
+            const int opposite_pos = i - cur_len - 1;
+            if (opposite_pos >= 0 and s.at(opposite_pos) == s.at(i)) return cursor;
+            cursor = nodes.at(cursor).suffix_link();
+        }
+    }
+
+    // 文字列 s を追加する。 Complexity: O(|s|)
+    // callback(i, cursor) は s[0:(i + 1)] が追加された後の nodes[cursor] に対して行う処理
+    template <class Callback> void add_string(const std::vector<Key> &s, Callback callback) {
+        int cursor = 1;
+
+        for (int i = 0; i < (int)s.size(); ++i) {
+
+            cursor = find_next_suffix(s, i, cursor);
+
+            int ch = nodes.at(cursor).get_child(s.at(i));
+
+            if (ch < 0) {
+                const int nxt_cursor = nodes.size();
+                const int new_length = nodes.at(cursor).length() + 2;
+
+                int new_suffix_link_par = find_next_suffix(s, i, nodes.at(cursor).suffix_link());
+                int new_suffix_link = nodes.at(new_suffix_link_par).get_child(s.at(i));
+                if (new_suffix_link < 0) new_suffix_link = 1;
+
+                nodes.at(cursor).set_child(s.at(i), nxt_cursor);
+                nodes.push_back(Node<Key>(new_suffix_link, new_length));
+                cursor = nxt_cursor;
+
+            } else {
+                cursor = ch;
+            }
+
+            callback(i, cursor);
+        }
+    }
+
+    template <class Callback> void add_string(const std::string &s, Callback callback) {
+        add_string(std::vector<Key>{s.cbegin(), s.cend()}, callback);
+    }
+
+    template <class Vec> void add_string(const Vec &s) {
+        add_string(s, [](int, int) {});
+    }
+};
+
+} // namespace palindromic_tree

string/palindromic_tree.md

@@ -0,0 +1,42 @@
+---
+title: Palindromic tree / eertree （回文木）
+documentation_of: ./palindromic_tree.hpp
+---
+
+文字列に現れる回文を効率的に管理するデータ構造．与えられた文字列 $S$ に対して $O(|S| \log \sigma)$ $(\sigma = |\Sigma|)$ で構築可能．
+
+## eertree とは
+
+- 各頂点が，与えられた文字列の（連続）部分文字列である distinct な回文に対応する．
+  - 例外として，長さ $-1$, $0$ の空文字列を表現するダミーの頂点がそれぞれ存在する．したがって頂点数は正確には回文の種類数 +2 となる．
+  - 頂点数は $|S| + 2$ 以下である．
+- 各頂点から，その回文の suffix である最長回文の頂点への辺 (suffix link) が生えている．
+  - 例外として，長さ $0$ の空文字列のダミー頂点の suffix link は長さ $-1$ の空文字列へ生えている．
+  - 長さ $-1$ の空文字列の suffix link は存在しない．つまり， suffix link を辺としたグラフはこの頂点を根とする根付き木となる．
+
+## 使用方法
+
+```cpp
+string S = "sakanakanandaka";
+palindromic_tree::Tree<char> tree;
+tree.add_string(S);
+
+// コールバック関数も使用可能．
+// 第一引数は S 上の index (0, ..., |S| - 1), 第二引数は tree.nodes の index.
+// 1 文字読む毎に tree 上のどこにいるか分かるので出現回数カウント等が可能．
+vector<int> dp(S.size() + 2);
+auto callback = [&](int str_idx, int node_idx) -> void { dp.at(node_idx)++; };
+
+tree.add_string(S, callback);
+```
+
+## 問題例
+
+- [No.263 Common Palindromes Extra - yukicoder](https://yukicoder.me/problems/no/263)
+- [No.2606 Mirror Relay - yukicoder](https://yukicoder.me/problems/no/2606)
+
+## 参考文献・リンク
+
+- [1] M. Rubinchik and A. M. Shur, "EERTREE: An efficient data structure for processing palindromes in strings," European Journal of Combinatorics 68, 249-265, 2018. [arXiv](https://arxiv.org/abs/1506.04862)
+- [Palindromic Tree - math314のブログ](https://math314.hateblo.jp/entry/2016/12/19/005919)
+- [すごい！EERTREE！いごす - 誤読](https://mojashi.hatenablog.com/entry/2017/07/17/155520)

string/test/palindromic_tree.yuki2606.test.cpp

@@ -0,0 +1,35 @@
+#define PROBLEM "https://yukicoder.me/problems/no/2606"
+#include "../palindromic_tree.hpp"
+
+#include <algorithm>
+#include <iostream>
+#include <string>
+#include <vector>
+
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    string S;
+    cin >> S;
+
+    palindromic_tree::Tree<char> tree;
+
+    vector<long long> visitcnt(S.size() + 2);
+    tree.add_string(S, [&](int, int node_idx) { visitcnt.at(node_idx)++; });
+
+    const int V = tree.nodes.size();
+    for (int i = V - 1; i > 0; --i) visitcnt.at(tree.nodes.at(i).suffix_link()) += visitcnt.at(i);
+
+    vector<vector<int>> children(V);
+    for (int i = 1; i < V; ++i) children.at(tree.nodes.at(i).suffix_link()).push_back(i);
+
+    vector<long long> dp(V, 0);
+    for (int i = 0; i < V; ++i) {
+        dp.at(i) += visitcnt.at(i) * max(tree.nodes.at(i).length(), 0);
+        for (int ch : children.at(i)) dp.at(ch) += dp.at(i);
+    }
+
+    cout << *max_element(dp.begin(), dp.end()) << '\n';
+}

string/test/palindromic_tree.yuki263.test.cpp

@@ -0,0 +1,37 @@
+#define PROBLEM "https://yukicoder.me/problems/no/263"
+#include "../palindromic_tree.hpp"
+
+#include <iostream>
+#include <string>
+#include <vector>
+
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    string S, T;
+    cin >> S >> T;
+
+    palindromic_tree::Tree<char> tree;
+
+    vector<long long> visitcnt(S.size() + T.size() + 2);
+
+    tree.add_string(S, [&](int, int node_idx) { visitcnt.at(node_idx)++; });
+    tree.add_string(T);
+
+    const int V = tree.nodes.size();
+    for (int v = V - 1; v > 0; --v) {
+        visitcnt.at(tree.nodes.at(v).suffix_link()) += visitcnt.at(v);
+    }
+
+    // 0 と 1 はダミーなのでカウントしてはいけない
+    visitcnt.at(0) = visitcnt.at(1) = 0;
+
+    for (int v = 1; v < V; ++v) visitcnt.at(v) += visitcnt.at(tree.nodes.at(v).suffix_link());
+
+    long long ret = 0;
+    tree.add_string(T, [&](int, int node_idx) { ret += visitcnt.at(node_idx); });
+
+    cout << ret << '\n';
+}