PalindromicTree.cpp

#include <iostream>
#include <vector>
#include <map>

using namespace std;

struct Node {
    map<char, int> next;    // nodeの両端にcharをつけたときにいくnodeのindex
    int len = 0;            // 回文の長さ
    int sufflink = 0;       // nodeのsuffix palindromeであるnodeのindex
    int num = 0;
};

class PalindromicTree {
public:
    string str;
    vector<Node> tree;
    int last_node_idx;                  // 一番最後にできたnodeのidx
    int max_suffix_palindrome_idx;      // 現在の最長の接尾辞回文のnodeのidx

    PalindromicTree(string s) : str(s) {
        this->tree.resize(s.size() + 2);
        this->last_node_idx = 1;
        this->max_suffix_palindrome_idx = 1;

        this->tree[0].len = -1;
        this->tree[0].sufflink = 0;
        this->tree[1].len = 0;
        this->tree[1].sufflink = 0;
    }

    bool addLetter(int pos) {
        
        char now_char = this->str[pos];

        // suffix linkをたどり，xAxを見つける
        int A = max_suffix_palindrome_idx;    // 現在着目してるnodeのidx
        while (true) {
            int curlen = tree[A].len; // Aの長さ
            if (pos - 1 - curlen >= 0 && this->str[pos - 1 - curlen] == now_char) {
                break;
            }
            A = tree[A].sufflink;
        }

        // すでにnodeがある
        if (this->tree[A].next[now_char] != 0) {
            max_suffix_palindrome_idx = tree[A].next[now_char];
            return false;
        }

        // ノードを追加
        max_suffix_palindrome_idx = ++this->last_node_idx;
        Node &new_node = tree[this->last_node_idx];
        new_node.len = tree[A].len + 2;
        tree[A].next[now_char] = this->last_node_idx;   // AからxAxにedgeを張る

        // 長さが1の回文のsuffix linkは固定
        if (new_node.len == 1) {
            new_node.sufflink = 1; // 1文字のnodeのsuffix link先は空文字のnode
            new_node.num = 1;
            return true;
        }

        // さらにsuffix linkをたどり，xBxを見つける
        int B = A;
        while (true) {
            B = tree[B].sufflink;
            int curlen = tree[B].len; // Bの長さ
            if (pos - 1 - curlen >= 0 && this->str[pos - 1 - curlen] == now_char) {
                break;
            }
        }

        new_node.sufflink = tree[B].next[now_char];   // xAxからxBxへsufifixを張る
        new_node.num = 1 + tree[new_node.sufflink].num;

        return true;
    }
};

/*
sに含まれる回文ののべ数をもとめる
aaaならaが3個，aaが2個，aaaが1個で合計6個の回文が含まれる
*/
int num_of_palindrome(string s) {
    PalindromicTree pt(s);
    int ans = 0;
    for (int i = 0; i < s.size(); i++) {
        pt.addLetter(i);
        ans += pt.tree[pt.max_suffix_palindrome_idx].num;
    }
    return ans;
}

int main() {
    string s = "bbcbcb";

    cout << num_of_palindrome(s) << endl;

    return 0;
}