查找字典中某个公共前缀的所有单词

//Trie.h
/*
 * trie树统计公共前缀的单词集合、个数
 */

#include <set>
#include <string>
using namespace std;

#define BRANCH 26

struct Node
{
	Node * next[BRANCH];
	int prefix; // 此前缀的单词个数,为验证结果而加
	bool isStr;
	Node():prefix(0), isStr(false)
	{
		memset(next, 0, sizeof(next));
	}
};

class Trie
{
private:
	Node * root;
	set<string> * preSet;

    /* private:
	 * DFS() and freeNextArr() are private functions
	 *
	/* DFS find all the words from Node "ptr"
	 * 实际就是寻找“以当前前缀本身为单词”的单词
	 */
	void DFS(Node * ptr, string preStr)
	{
		for(int i = 0; i < BRANCH; ++i)
		{
			if(ptr->next[i] != NULL)
			{
				char chr[] = {'a' + i, '\0'};
				DFS(ptr->next[i],preStr+string(chr));
			}
			if(ptr->isStr)  // 存在当前前缀的单词
			{
				this->preSet->insert(preStr);
			}
		}
	}

	// free the next array
	void freeNextArr(Node * ptr)
	{
		for(int i = 0; i < BRANCH; ++i)
		{
			if(ptr->next[i] != NULL)
			{
				freeNextArr(ptr->next[i]);
			}
		}
		delete ptr;
	}

public:
	Trie()
	{
		root = new Node();
		preSet = new set<string> ();
	}

	// insert a word
	void insert(const char * word)
	{
		if(NULL == word)
		{
			return ;
		}
		Node * ptr = root;
		while(*word)
		{
			if(ptr->next[*word-'a'] == NULL)
			{
				ptr->next[*word-'a'] = new Node();
			}
			ptr->prefix ++;
			ptr = ptr->next[*word-'a'];
			++word;
		}
		ptr->isStr = true;
	}

	// find the location of the word
	Node * search(const char * word)
	{
		if(NULL == word)
		{
			return NULL;
		}
		Node * ptr = root;
		while(ptr && *word)
		{
			ptr = ptr->next[*word-'a'];
			++word;
		}
		return ptr;
	}

	// find the common prefix words set
	set<string> * findCommonPrefix_Set(const char * prefix)
	{
		if(NULL == prefix)
		{
			return NULL;
		}
		Node * ptr = search(prefix); // find location of prefix
		if(ptr == NULL)              // prefix not exist
		{
			return NULL;
		}
		DFS(ptr, string(prefix));   // DFS find commonPre words

		return this->preSet;
	}

	// clear the preSet
	void clearPreSet()
	{
		this->preSet->clear();
	}

	// free the memory
	void freeMemory()
	{
		this->freeNextArr(root);
		delete this->preSet;
	}
};



测试程序：

#include <iostream>
#include "Trie.h"

void main()
{
	Trie comPreTrie;
	comPreTrie.insert("abc");
	comPreTrie.insert("abcd");
	comPreTrie.insert("abcde");
	comPreTrie.insert("abcdef");
	comPreTrie.insert("abcdefg");
	comPreTrie.insert("cde");

	/* test the common prefix "abcd"
	 * 注，测试时默认输入合法字符，即 a-z
	 *    本文程序不对此做出检查
	 */
	Node * ptr = comPreTrie.search("abcd");
	if(ptr)
	{
		printf("num of words prefix by \"abcd\" is %d\n",
			ptr->isStr ? ptr->prefix + 1 : ptr->prefix);
	}else
	{
		printf("ptr for \"abcd\" is null\n");
	}

	set<string> * comPreSet
		= comPreTrie.findCommonPrefix_Set("abcd");
	if(comPreSet)
	{
		for(set<string>::iterator it = comPreSet->begin();
			it != comPreSet->end(); ++it)
		{
			printf("%s\n", (*it).c_str());
		}
	}else
	{
		printf("comPreSet for \"abcd\" is null\n");
	}
	comPreTrie.clearPreSet();

	printf("---------------\n");

	/* test the common prefix "word" */
	ptr = comPreTrie.search("word");
	if(ptr)
	{
		printf("num of words prefix by \"word\" is %d\n",
			ptr->isStr ? ptr->prefix + 1 : ptr->prefix);
	}else
	{
		printf("ptr for \"word\" is null\n");
	}

	comPreSet = comPreTrie.findCommonPrefix_Set("word");
	if(comPreSet)
	{
		for(set<string>::iterator it = comPreSet->begin();
			it != comPreSet->end(); ++it)
		{
			printf("%s\n", (*it).c_str());
		}
	}else
	{
		printf("comPreSet for \"word\" is null\n");
	}
	comPreTrie.clearPreSet();

	// free memory
	comPreTrie.freeMemory();
}