3135. 홍준이의 사전놀이

LEVEL D5

 

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV_6pTXqsXUDFAWS&categoryId=AV_6pTXqsXUDFAWS&categoryType=CODE

SW Expert Academy

 

트라이 자료구조를 공부하고 풀었다.

 

현재까지 저장 된 문자열 중 특정 문자열로 시작되는 문자열의 갯수를 찾는 문제였다.

 

INSERT가 들어오면 문자열을 저장하는 것이고, QUERY가 올 경우 저장 된 문자열 중 특정 문자열로 시작되는 문자열 갯수를 리턴하면 된다.

 

처음엔 QUERY가 들어 올 경우 저장 된 트리를 모두 찾아보는 방식으로 풀었으나 시간초과가 발생하였다.

 

소스코드

#include<stdio.h>
#include<string.h>
struct trie {
	trie* next[26];
	bool finish;
	trie() : finish(false) {
		memset(next, 0, sizeof(next));
	}
	~trie() {
		for (int i = 0; i < 26; i++) {
			if (next[i]) {
				delete next[i];
			}
		}
	}

	void insert(const char *key) {
		if (*key == '\0') {
			finish = true;
		}
		else {
			int cur = *key - 'a';
			if (next[cur] == NULL) {
				next[cur] = new trie();
			}
			next[cur]->insert(key + 1);
		}
	}

	trie* find(const char *key) {
		if (*key == '\0') return this;
		int cur = *key - 'a';
		if (next[cur] == NULL) {
			return 0;
		}
		return next[cur]->find(key + 1);
	}

	int findNodeCount(trie *cur)
	{
		int count = 0;
		if (cur->finish)
			count++;
		for (int i = 0; i < 26; i++) {
			if (cur->next[i] == 0) continue;
			count += findNodeCount(cur->next[i]);
		}
		return count;
	}

	int query(const char *key) {
		trie* cur = find(key);
		if (cur) {
			int sum = 0;
			if (cur->finish)
				sum++;
			for (int i = 0; i < 26; i++) {
				if (cur->next[i] == NULL) continue;
				sum += findNodeCount(cur->next[i]);
			}
			return sum;
		}
		else {
			return 0;
		}
	}
};
trie *root;
void init(void) {
    root = new trie;
}

void insert(int buffer_size, char *buf) {
    root->insert(buf);
}

int query(int buffer_size, char *buf) {
	return root->query(buf);
}

   

 

 

두 번째 생각한 방법은 INSERT 할 경우 거쳐 가는 노드의 childCount값을 증가시켜 자식의 갯수를 체크하도록 한 것이다.

 

이 방법을 사용하면 특정 노드의 INSERT시 하위 노드의 카운트가 증가하므로 QUERY시 특정 문자열의 노드를 찾은 후 해당 노드의 childCount값이 바로 내 아래 있는 노드의 갯수가 되므로 좀 더 빠른 방법이 되었다.

 

#include<stdio.h>
#include<string.h>
struct trie {
	trie* next[26];
	bool finish;
	int childCount;
	trie() : finish(false), childCount(0) {
		memset(next, 0, sizeof(next));
	}
	~trie() {
		for (int i = 0; i < 26; i++) {
			if (next[i]) {
				delete next[i];
			}
		}
	}

	void insert(const char *key) {
		if (*key == '\0') {
			finish = true;
			childCount++;
		}
		else {
			int cur = *key - 'a';
			if (next[cur] == NULL) {
				next[cur] = new trie();
			}
			next[cur]->insert(key + 1);
			childCount++;
		}
	}

	trie* find(const char *key) {
		if (*key == '\0') return this;
		int cur = *key - 'a';
		if (next[cur] == NULL) {
			return 0;
		}
		return next[cur]->find(key + 1);
	}

	int query(const char *key) {
		trie* cur = find(key);
		if (cur) {
			return cur->childCount;
		}
		else {
			return 0;
		}
	}
};

trie *root;
void init(void) {
    root = new trie;
}

void insert(int buffer_size, char *buf) {
    root->insert(buf);
}

int query(int buffer_size, char *buf) {
	return root->query(buf);
}