Good Substrings CF271D

SAM上的DP

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <queue>
#include <algorithm>
#include <vector>
#include <cstring>
#include <stack>
#include <cctype>
#include <utility>   
#include <map>
#include <string>  
#include <climits> 
#include <set>
#include <string>    
#include <sstream>
#include <utility>   
#include <ctime>

using std::priority_queue;
using std::vector;
using std::swap;
using std::stack;
using std::sort;
using std::max;
using std::min;
using std::pair;
using std::map;
using std::string;
using std::cin;
using std::cout;
using std::set;
using std::queue;
using std::string;
using std::stringstream;
using std::make_pair;
using std::getline;
using std::greater;
using std::endl;
using std::multimap;
using std::deque;

typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> PAIR;
typedef multimap<int, int> MMAP;

const int MAXN(1510);
const int MAXM(5010);
const int MAXE(10010);
const int HSIZE(13131);
const int SIGMA_SIZE(26);
const int MAXH(19);
const int INFI((INT_MAX-1) >> 1);
const ULL BASE(31);
const LL LIM(10000000);
const int INV(-10000);

int K;
int score[26];

struct SAM
{
    struct NODE
    {
        int len;
        NODE *f, *ch[SIGMA_SIZE];
        int table[MAXN];
    };
    NODE *root, *last;
    NODE pool[MAXN << 1];
    int size;
    void init()
    {
        root = last = pool;
        root->f = 0;
        root->len = 0;
        memset(root->ch, 0, sizeof(root->ch));
        memset(root->table, -1, sizeof(root->table[0])*(K+1));
        size = 1;
    }

    NODE *newnode(int tl)
    {
        pool[size].len = tl;
        memset(pool[size].ch, 0, sizeof(pool[size].ch));
        memset(pool[size].table, -1, sizeof(pool[size].table[0])*(K+1));
        return pool+size++;
    }
    void extend(int id)
    {
        NODE *p = last, *np = newnode(last->len+1);
        last = np;
        while(p && p->ch[id] == 0)
            p->ch[id] = np, p = p->f;
        if(p == 0)
            np->f = root;
        else
        {
            NODE *q = p->ch[id];
            if(p->len+1 == q->len)
                np->f = q;
            else
            {
                NODE *nq = newnode(p->len+1);
                memcpy(nq->ch, q->ch, sizeof(nq->ch));
                nq->f = q->f;
                q->f = np->f = nq;
                while(p && p->ch[id] == q)
                    p->ch[id] = nq, p = p->f;
            }
        }
    }
    int dfs(NODE *p, int cnt)
    {
        if(cnt > K)
            return 0;
        if(p->table[cnt] != -1)
            return p->table[cnt];
        int &cur = p->table[cnt];
        cur = 1;
        for(int i = 0; i < SIGMA_SIZE; ++i)
            if(p->ch[i])
                cur += dfs(p->ch[i], cnt+score[i]);
        return cur;
    }
};

SAM sam;
char str[MAXN];

int main()
{
    while(~scanf("%s", str))
    {
        sam.init();
        char temp;
        for(int i = 0; i < 26; ++i)
        {
            scanf(" %c", &temp);
            score[i] = (temp-'0')^1;
        }
        scanf("%d", &K);
        for(char *sp = str; *sp; ++sp)
            sam.extend(*sp-'a');
        sam.dfs(sam.root, 0);
        printf("%d\n", sam.root->table[0]-1);
    }
    return 0;
}