UVA 1519 - Dictionary Size(Trie树)

UVA 1519 - Dictionary Size

题目链接

题意：有一个字典，里面包含一些词，要求组合新词，新词必须来自原字典，或者由原字典的字符串的非空前缀和非空后缀组成，问一共能组成多少个新词

思路：建Trie树，可以求出不同的前缀和后缀个数，然后相乘，这样做会有一部分重复的
比如Aaaa，aaaA的情况，就重复了，去重的方法可以推理出来
假设前缀A后面有x个a，后缀y个a后面一个A，那么这x和y个一共一边各有(x+1)和(y+1)种情况，相乘一共是(x+1)(y+1)种，那么实际上一共只能组成1 + x + y(比如上面的例子，一共就能组成空串,a, aa, aaa, aaaa, aaaaa, aaaaaa种)这样减掉就得到x*y,这样一来只要在建Trie的过程中，记录下每种字符各出现多少个就能算了(注意开头的单独一个不能算，因为要非空)，然后还有一点要注意的，就是长度为1的字符串，由于题目单词可以来自原字典，而长度1的字符串在计算过程并没有算进去，所以还要单独出来算

代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int MAXNODE = 400005;
const int SIGMA_SIZE = 26;

struct Trie {
    int ch[MAXNODE][SIGMA_SIZE];
    int cnt[SIGMA_SIZE];
    int sz;

    void init() {
	sz = 1;
	memset(ch[0], 0, sizeof(ch[0]));
	memset(cnt, 0, sizeof(cnt));
    }

    int idx(char c) {
	return c - 'a';
    }

    void insert(char *s) {
	int n = strlen(s);
	int u = 0, v;
	for (int i = 0; i < n; i++) {
	    int c = idx(s[i]);
	    if (!ch[u][c]) {
		memset(ch[sz], 0, sizeof(ch[sz]));
		ch[u][c] = sz++;
		if (i) cnt[c]++;
	    }
	    u = ch[u][c];
	}
    }
}pre, suf;

typedef long long ll;
int n, vis[26];
char str[45];

int main() {
    while (~scanf("%d", &n)) {
	pre.init(); suf.init();
	memset(vis, 0, sizeof(vis));
	while (n--) {
	    scanf("%s", str);
	    if (strlen(str) == 1)
		vis[str[0] - 'a'] = 1;
	    pre.insert(str);
	    reverse(str, str + strlen(str));
	    suf.insert(str);
	}
	ll ans = (ll)(pre.sz - 1) * (suf.sz - 1);
	for (int i = 0; i < 26; i++)
	    ans -= (ll)pre.cnt[i] * suf.cnt[i] - vis[i];
	printf("%lld\n", ans);
    }
    return 0;
}