一个Trie树的实现

Trie树实现代码：

// Max branch number. Such as 26 for lower case letter set, 10 for digit character set.
const int MAX_BRANCH_NUM = 26;

// Node of Trie Tree        
class CTrieNode
{
public:
	CTrieNode()
	{
		nPrefixCount = 0;
		nWordCount = 0;
		memset(next, NULL, sizeof(next));
	}
public:
	int nPrefixCount;    // Count of prefix in the Trie Tree
	int nWordCount;      // Count of word in the Trie Tree
	CTrieNode* next[MAX_BRANCH_NUM];  // Point to all the nodes in the next level
};

// Trie Tree class
class CTrieTree
{
public:
	CTrieTree();    // Constructor
	~CTrieTree();   // Destructor
	int GetIndex(char ch) const;           // Get the index of 'ch' in the 'next' array
	bool CheckWord(const char *s) const;   // Check if there is illegal character
	bool InsertWord(const char *s) const;  // Insert a word
	int SearchPrefix(const char *s) const; // Search the count of a prefix
	int SearchWord(const char *s) const;   // Search the count of a word
	bool DeleteWord(const char *s) const;  // Delete a word
private:
	void FreeTreeNodes(CTrieNode * root);  // Free those tree nodes
private:
	CTrieNode *pRoot; // Root node pointer of Trie Tree
};

// Constructor
CTrieTree::CTrieTree()
{
	pRoot = new CTrieNode; 
}

// Destructor
CTrieTree::~CTrieTree()
{
	FreeTreeNodes(pRoot);
}

// Get the index of 'ch' in the 'next' array
int CTrieTree::GetIndex(char ch) const
{
	// This is for lower case letter set
	return ch - 'a';
}

// Check if there is illegal character in the word
bool CTrieTree::CheckWord(const char *s) const
{
	// This is for lower case letter set
	for(int i = 0; s[i] != '\0'; i++)
		if(s[i] < 'a' || s[i] > 'z')
			return false;
	return true;
}

// Insert a word
bool CTrieTree::InsertWord(const char *s) const
{
	// Can't insert a null string or an word containing illegal character
	if((s[0] == '\0') || (CheckWord(s)== false))
		return false;
	CTrieNode * pCurrent = pRoot;
	int i = 0;
	while(1)
	{
		int nIndex = GetIndex(s[i]);
		if(pCurrent->next[nIndex] == NULL)
			pCurrent->next[nIndex] = new CTrieNode;
		pCurrent = pCurrent->next[nIndex];
		pCurrent->nPrefixCount++;
		i++;
		if(s[i] == '\0')
		{
			pCurrent->nWordCount++;
			break;
		}
	}
	return true;
};

// Search the count of a prefix in the Trie Tree
int CTrieTree::SearchPrefix(const char *s) const
{
	// There are no null prefix or a prefix containing illegal character
	if((s[0] == '\0') || (CheckWord(s)== false))
		return 0;
	CTrieNode * pCurrent = pRoot;
	int i = 0;
	while(1)
	{
		int nIndex = GetIndex(s[i]);
		pCurrent = pCurrent->next[nIndex];
		if((pCurrent == NULL) || (pCurrent->nPrefixCount == 0))
			return 0;
		i++;
		if(s[i] == '\0')
			return pCurrent->nPrefixCount;
	}
}

// Search the count of a word in the Trie Tree
int CTrieTree::SearchWord(const char *s) const
{
	// There are no null word or a word containing illegal character
	if((s[0] == '\0') || (CheckWord(s)== false))
		return 0;
	CTrieNode * pCurrent = pRoot;
	int i = 0;
	while(1)
	{
		int nIndex = GetIndex(s[i]);
		pCurrent = pCurrent->next[nIndex];
		if((pCurrent == NULL) || (pCurrent->nPrefixCount == 0))
			return 0;
		i++;
		if(s[i] == '\0')
			return pCurrent->nWordCount;
	}
}

// Delete a word
bool CTrieTree::DeleteWord(const char *s) const
{
	if(!SearchWord(s))
		return false;
	CTrieNode * pCurrent = pRoot;
	int i = 0;
	while(1)
	{
		int nIndex = GetIndex(s[i]);
		pCurrent = pCurrent->next[nIndex];
		pCurrent->nPrefixCount--;
		i++;
		if(s[i] == '\0')
		{
			pCurrent->nWordCount--;
			break;
		}
	}
	return true;
}

// Free those tree nodes
void CTrieTree::FreeTreeNodes(CTrieNode * pRoot)
{
	if(pRoot == NULL)
		return;
	for(int i = 0; i < MAX_BRANCH_NUM; i++)
		FreeTreeNodes(pRoot->next[i]);
	free(pRoot);
	pRoot = NULL;
}