Codeforces 557E Ann and Half-Palindrome DP+字典树

求字典序第K大的半回文串，半回文串满足第一位和最后一位相等，第三位和倒数第三位相等。。。。 先用dp求出所有的半回文串，然后插入到字典树中统计个数，根据个数求第K大。

#include 
 
 
  
  
#include 
  
  
   
   
char s[5010];
int k;
int dp[5010][5010];
int ch[5010*5010/2][2];
int sum[5010*5010/2];
int val[5010*5010/2];
int sz;

void init() {
    sz = 1;
}

void dfs(int u, int x, int n, int pos) {
    if(n == x) {
        return;
    }
    int c = s[x]-'a';

    if(ch[u][c] == 0)
    {
        ch[u][c] = sz++;
    }
    int v = ch[u][c];
    dfs(v, x+1, n, pos);
    if(dp[pos][x] == 1) {
        val[v] += 1;
    }
    sum[v] = sum[ch[v][0]] + sum[ch[v][1]] + val[v];
}

void insert(int pos) {

    int u = 0, n = strlen(s);
    dfs(u, pos, n, pos);
}
void find(int k) {
    char ans[5010];
    int t = 0;
    int u = 0;
    while(k > 0) {
        int c = ch[u][0];
        //System.out.println(sum[u] + " " + sum[ch[u][0]] + " " + sum[ch[u][1]] + " " + k);
        int num = sum[u]-(sum[ch[u][0]]+sum[ch[u][1]]);
        if(u == 0)
            num = 0;
        if(k <= num)
            break;
        k -= num;
        if(sum[c] >= k) {
            u = c;
            ans[t++] = 'a';
        }
        else {
            k -= sum[c];
            u = ch[u][1];
            ans[t++] = 'b';
        }

    }
    ans[t] = 0;
    printf("%s", ans);
}
void dfs2(int u) {
    if(u == 0)
        return;
    dfs2(ch[u][0]);
    dfs2(ch[u][1]);
    sum[u] += sum[ch[u][0]] + sum[ch[u][1]];
}
int main() {

    scanf("%s %d", s, &k);
    int len = strlen(s);
    for(int i = 0; i < len; i++) {
        dp[i][i] = 1;
    }
    for(int i = 0 ; i+1 < len; i++) {
        if(s[i] == s[i+1])
            dp[i][i+1] = 1;
    }
    for(int i = 0 ; i+2 < len; i++) {
        if(s[i] == s[i+2])
            dp[i][i+2] = 1;
    }
    for(int i = 0 ; i+3 < len; i++) {
        if(s[i] == s[i+3])
            dp[i][i+3] = 1;
    }

    for(int i = 5; i <= len; i++) {
        for(int j = 0; j+i <= len; j++) {
            if(s[j] == s[j+i-1] && dp[j+2][j+i-3] == 1) {
                dp[j][j+i-1] = 1;
            }
        }
    }
    init();
    for(int i = 0; i < len; i++) {
        insert(i);
    }
    //dfs2(ch[0][0]);
    //dfs2(ch[0][1]);
    find(k);
}