557E Anya and Half-palindrome

editorial を見て解くも TLE しまくった問題。結局他の人コードを参考に Trie の add 関数を修正して解いた。
コメントにあるように int と size_t を比較して意図どおりに動かないというバグに見舞われた。
class Div2E {
public:
    struct TrieNode {
        std::array<TrieNode*, 2> child;
        int      cnt;
        TrieNode() : child({nullptr,nullptr}), cnt(0) {}
        ~TrieNode() {
                for (int i = 0; i < 2; ++i)
                    delete child[i];
        }
    };
    vector<vector<bool>> isHP;
    // O(N)
    // 部分文字列 str[start...cur+1] をたどる
    void add(TrieNode *t, const string &str, int start, int cur) {
            assert(start < str.length());
            // str.length が size_t なので cur >= str.length
            // だけだと -1 >= str.lengt() になってしまう
            if (cur >= (int)str.length())
                return;
            // 部分文字列 str[start...cur+1] が half palindrome ならノードをカウント
            // 重複を許していることに注意
            if (start <= cur && isHP[start][cur+1])
                ++t->cnt;
            // 次のノードへ
            int idx = (str[cur+1]=='a') ? 0 : 1;
            if (!t->child[idx])
                t->child[idx] = new TrieNode();
            add(t->child[idx], str, start, cur + 1);
    }
    int  count(TrieNode *t)  { return t ? t->cnt : 0; }
    void update(TrieNode *t) {
            if (!t)
                return;
            update(t->child[0]);
            update(t->child[1]);
            t->cnt += count(t->child[0]) + count(t->child[1]);
    }
    TrieNode* next(TrieNode *cur, char c) {
            if (!cur)
                return nullptr;
            return (c == 'a')? cur->child[0] : cur->child[1];
    }
    void solve(void) {
            string str;
            int K;
            cin>>str>>K;

            int N = str.length();

            // まず half-palindrome な部分文字列を探す
            // isHP[i][j] := str[i..j] が half palindrome か
            isHP.resize(N+1, vector<bool>(N+1, false));

            // O(N^2)
            FOR(w, 1, N+1)
            for (int i = 0; i+w <= N; ++i)
            {
                if (str[i] == str[i+w-1] && (w <= 4 || isHP[i+2][(i+w)-2]))
                    isHP[i][i+w] = true;
            }
            TrieNode root;

            // Trie に結果を挿入
            for (int i = 0; i < N; ++i)
                add(&root, str, i, i-1); // root node を "" として扱うように cur=i-1 で呼ぶ

            // cnt を更新
            update(&root);

            // 順番から再現する
            string ans = "";
            TrieNode *cur = &root;
            int k = count(cur);
            REP(i, N)
            {
                auto ta = next(cur, 'a');
                auto tb = next(cur, 'b');
                int a = count(ta);
                int b = count(tb);

                int l = k-(a+b);
                K -= l;
                if (K <= 0)
                    break;
                if (K <= a)
                {
                    ans += "a";
                    cur = ta;
                    k = a;
                }
                else
                {
                    K -= a;
                    ans += "b";
                    cur = tb;
                    k = b;
                }
            }
            cout<<ans<<endl;
    }
};