Trie

트리가 아니라 트라이

2018-08-23

태그: dev

문자열을 비교하고 검색하는 알고리즘은 항상 중요하다. 뭐 여러가지 알고리즘이 많은데 요즘 가장 재밌게 구현한 알고리즘은 단연 트라이 다(트리가 아니다!). TRIE는 Edward Fredkin이 만들었고 reTRIEval이라는 단어에서 따왔다고 한다. 태생이 검색을 위한 놈이다. 간단히 설명하면 트리 모양을 한 결정적 유한 오토마타다.

문자열을 키로 어떤 데이터를 관리해야 하는 상황에서 데이터의 추가/ 검색이 잦은 경우 고려해볼만 한 자료구조다. 대신 문자(character) 하나에 트리의 엣지 하나가 대응되기 때문에, 메모리를 많이 못쓰는 경우에는 쓸 수 없다. 적당히 데이터의 키가 소문자로만 이뤄진 경우, 그러니까 기껏해야 26자만 필요한 경우에는 꽤 유용하게 쓸 수 있다.

트라이로 문자열 관련 문제를 풀면 해쉬 같은 방법에 비해 속도가 엄청나게 빨라진다. 해쉬는 항상 충돌의 위험이 있어서 해쉬 값이 같으면 더블 체크를 해줘야 하기 때문에.

트라이로 풀 수 있는 문제 중에 재밌는 건 간단한 디렉토리 커맨드 구현이다. 리눅스의 ls, mkdir, rm, cd 키워드를 구현해보자.

일단 두 가지 자료구조를 쓸거다. 하나는 트라이고 하나는 유사 inode이다.

struct INode;
struct Trie{
  Trie *children[26]; // lower-case only
  int count; // count of reference
  bool valid; // for fast remove
  INode *dir; // only valid when terminal
};

트라이는 간단하다. 일단 소문자만 사용하기 때문에 각 엣지(소문자)에 대응할 26개의 자식 포인터가 있다. count 는 참조 횟수로, 예를 들어 현재 디렉토리에 a와 ab라는 디렉토리 두개가 있으면, a 엣지의 카운트는 2이고 a 노드의 b 엣지 카운트는 1이 되게끔 할 거다. valid는 빠른 삭제를 위한 플래그이고, dir는 실제로 트라이의 끝에 매달릴 inode 포인터이다.

struct INode{
  INode *parent;
  Trie *sub;
  int size;
};

유사 아이노드는 실제 디렉토리 정보를 담을 자료구조다. cd .. 커맨드를 위해서 부모 디렉토리의 포인터도 하나 갖고 있고, 현재 디렉토리의 자식 디렉토리를 관리하기 위한 트라이 sub도 갖고 있다. 그리고 size는 자식 디렉토리의 개수가 아니라, 자식 디렉토리의 개수와 그 각각의 자식 디렉토리 안의 모든 자식 디렉토리 개수를 다 합친 값이다.

이렇게하면 일단 디렉토리 중복체크를 다음과 같이 O(1)만에 확인할 수 있다.

bool conflict(Trie *cur, char *path){
  Trie *trie = cur;
  int idx = 0;
  while(path[idx] != '\0'){
    int t_idx = path[idx] - 'a';
	if(trie->children[t_idx] == nullptr)
	  return false;
	trie = trie->children[t_idx];
	if(!trie->valid)
	  return false;
	++idx;
  }
  return trie->dir != nullptr;
}

쉽다. 중복체크할 디렉토리 이름의 문자열 하나하나를 가지고 현재 디렉토리의 sub(트라이)를 쭈욱 따라서 하나하나 확인해보면 된다. 엣지가 없거나 유효하지 않으면 false(노중복)다. 완전히 이름이 같은 경우 해당 트라이에 매달린 디렉토리가 있는지를(trie->dir이 널인지) 보면 된다.

이 중복체크 함수를 가지고 mkdir 함수를 만들면 된다.

bool mkdir(char *path){
  if(conflict(wd->sub, path)) return false;
  else{
    Trie *trie = wd->sub;
	int idx = 0;
	while(path[idx] != '\0'){
	  int t_idx = path[idx] - 'a';
	  if(trie->children[t_idx] == nullptr)
	    trie->children[t_idx] = trie_pool.alloca();
	  trie = trie->children[t_idx];
	  if(!trie->valid) trie->clear();
	  trie->count += 1;
	  ++idx;
	}
	trie->dir = inode_pool.alloca();
	trie->dir->parent = wd;
	propagate_inode_count(wd, 1);
	return true;
  }
}

먼저 중복체크를 하고, 중복이 없는 경우 디렉토리 문자열을 가지고 하나하나 따라가면서 엣지가 없으면 만들고, 삭제됐으면 초기화하고, 있으면 레퍼런스를 1씩 늘려준다. 최종적으로 디렉토리 이름을 다 따라갔으면 진짜 디렉토리를 뜻하는 아이노드를 만들어서 매달아주고 해당 아이노드의 부모 디렉토리를 지금의 작업 디렉토리(wd; working directory)로 만든다. 마지막으로 ls와 rm 커맨드를 위해 현재로부터 루트까지 모든 부모 디렉토리에게 개수 증가 1을 전파한다. 전파 함수는 아래와 같이 매우 쉽다.

void propagate_inode_count(INode *inode, int acc){
  if(inode == nullptr) return;
  inode->size += acc;
  propagate_inode_count(inode->parent, acc);
}

작업 디렉토리를 바꾸는 cd 함수는 다음과 같이 짜면 된다.

bool cd(char *path){
  if(path == "/"){
    wd = root;
	return true;
  }else if (path == ".."){
    if(wd == root) return false;
	else{
	  wd = wd->parent;
	  return true;
	}
  }else{
    Trie *trie = wd->sub;
	int idx = 0;
	while(path[idx] != '\t'){
	  int t_idx = path[idx] - 'a';
	  if(trie->children[t_idx] == nullptr) return false;
	  trie = trie->children[t_idx];
	  if(!trie->valid) return false;
	  ++idx;
	}

	while(trie != nullptr && trie->dir == nullptr){
	  Trie *next = nullptr;
	  for(int i=0; i<26; ++i)
	    if(trie->children[i] != nullptr && trie->children[i]->valid){
		  next = trie->children[i];
		  break;
		}
	  trie = next;
	}

	if(trie == nullptr || trie->dir == nullptr) return false;

    wd = trie->dir;
	return true;
  }
}

역시 트라이는 명확한 자료구조라 쉽게 구현할 수 있다. cd의 경우 탭 키가 있으면 알파벳 순으로 가장 빠른 자식 디렉토리로 가도록 하기 위해서 추가적인 검색을 더 했다.

모든 자식 디렉토리 및 자식 디렉토리의 리프 디렉토리의 개수를 반환하기 위한 ls는 다음과 같다.

int ls(char *path){
  if(path == "*") return wd->size;
  Trie *trie = wd->sub;
  int idx = 0;
  while(path[idx] != '*'){
    int t_idx = path[idx] - 'a';
	if(trie->children[t_idx] == nullptr) return 0;
	trie = trie->children[t_idx];
	if(!trie->valid) return 0;
	++idx;
  }

  int ls_cnt = 0;
  ls_cnt += trie->count;
  retrieve_all_subdir(trie, ls_cnt);
  return ls_cnt;
}

역시 비슷하게 디렉토리 이름을 가지고 트라이를 하나하나 따라가고(이거 자꾸 비슷한 funcionality가 중복되는 기분이 드는데 어떻게하면 깔끔하게 중복을 제거할 수 있을지 좀더 고민해봐야겠다), *을 만나면 이제 거기서부터 최종 개수를 센다. 먼저 * 이전까지의 문자열의 참조 개수, 즉 현재 디렉토리가 갖고 있는 자식 디렉토리의 개수를 세고, 거기다 해당 이름으로 시작하는 모~~~든 자식 디렉토리의 개수를 누적하면 된다. 우리는 아이노드의 size에 이를 잘 캐싱해뒀기 때문에 다음과 같이 하면 쉽게 구할 수 있다!

void retrieve_all_subdir(Trie *sub, int &acc){
  if(sub->dir != nullptr) acc += sub->dir->size;
  for(int i=0; i<26; ++i)
    if(sub->children[i] != nullptr && sub->children[i]->valid)
	  retrieve_all_subdir(sub->children[i], acc);
}

이제 남은 건 디렉토리를 삭제하고 삭제 개수를 반환하는 rm이다.

int rm(char *path){
  if(path == "*"){
    for(int i=0; i<26; ++i){
	  if(wd->sub->children[i] != nullptr)
	    wd->sub->children[i]->clear();
	}
	propagate_inode_count(wd, -(wd->size));
	int rm_cnt = wd->size;
	wd->size = 0;
	return rm_cnt;
  }

  Trie *trie = wd->sub;
  int idx = 0;
  while(path[idx] != '*'){
    int t_idx = path[idx] - 'a';
	if(trie->children[t_idx] == nullptr) return 0;
	trie = trie->children[t_idx];
	if(!trie->valid) return 0;
	++idx;
  }

  int cur_dir_rm_cnt = trie->count;
  int all_subdir_rm_cnt = cur_dir_rm_cnt;
  retrieve_all_subdir(trie, all_subdir_rm_cnt);

  trie->valid = false;

  idx = 0;
  trie = wd->sub;
  while(path[idx] != '*'){
    int t_idx = path[idx] - 'a';
	trie = trie->children[t_idx];
	trie->count -= cur_dir_rm_cnt;
	if(trie->count <= 0) trie->valid = false;
	++idx;
  }

  propagate_inode_count(wd, -all_subdir_rm_cnt);
  return all_subdir_rm_cnt;
}

이때까지 만든 함수들을 적절히 이용해서 짜깁기하면 쉽다. 삭제할 때 (1) 트라이 및 디렉토리 삭제, (2) 삭제한 개수 캐싱 이 두 가지만 제대로 하면 된다. 다 삭제하는 경우는 쉬우니 검색하는 경우를 보면 일단 원하는 트라이까지 따라오고 나면 해당 트라이의 참조 회수가 곧 현재 디렉토리에서 지워야하는 서브 디렉토리의 개수이다. 일단 이걸 캐싱하고, 나머지는 이 서브 디렉토리가 갖고 있는 모든 서브 디렉토리의 개수를 갖고 와서 빼야 한다. ls 만들 때 썼던 함수를 갖다 쓰면 쉽게 해결된다.