洛谷P2580 于是他错误的点名开始了(trie树模板题)

原题链接
本题属于trie树模板题，详见代码
评测时长

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <string>
using namespace std;
#define maxnode 1000050
#define sigma_size 30
struct Trie
{
	int sz;
	int ch[maxnode][sigma_size];
	int val[maxnode];
}tree;
int inite()
{
	tree.sz = 1;
	memset(tree.ch[0], 0, sizeof(tree.ch[0]));
}

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

void insert_(char* s)
{
	int u = 0,n = strlen(s);
	for (int i = 0; i < n; i++)
	{
		int c = idx(s[i]);
		if (!tree.ch[u][c])
		{
			memset(tree.ch[tree.sz], 0, sizeof(tree.ch[tree.sz]));
			tree.val[tree.sz] = 0;
			tree.ch[u][c] = tree.sz++;
		}
		u = tree.ch[u][c];
	}
	tree.val[u] = 1;
}
int query(char* s)
{
	int u = 0, n = strlen(s);
	for (int i = 0; i < n; i++)
	{
		int c = idx(s[i]);
		if (!tree.ch[u][c])
		{
			return 0;
		}
		u = tree.ch[u][c];
	}
	if (tree.val[u] == 0)
	{
		return 0;
	}
	if (tree.val[u] == 1)
	{
		tree.val[u] = 2;
		return 1;
	}
	if (tree.val[u] == 2)
	{
		return 2;
	}
}
int main()
{
	int n; char a[1000];
	scanf("%d", &n);
	inite();
	while (n--)
	{
		scanf("%s", a);
		insert_(a);
	}
	int m;
	scanf("%d", &m);
	while (m--)
	{
		scanf("%s", a);
		int k = query(a);
		if (k == 0)printf("WRONG\n");
		if (k == 1)printf("OK\n");
		if (k == 2)printf("REPEAT\n");
	}
}