C言語で簡単なTrieを実装した

コーディングテストを受けた。要件としては文字列の探索だったので「Trieを使うんだろうな〜」という知識はあったがその場で実装できなかったのでいま筆を取っている。*1*2

[https://ja.wikipedia.org/wiki/(データ構造)]

https://ja.wikipedia.org/wiki/トライ(データ構造) から拝借

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

typedef struct s_trie_node {
    struct s_trie_node *children[255];
    bool is_end_of_word;
} t_trie_node;

t_trie_node *new_node(void) {
    t_trie_node *node = (t_trie_node *)malloc(sizeof(t_trie_node));
    for (int i = 0; i < 255; i++)
        node->children[i] = NULL;
    node->is_end_of_word = false;
    return node;
}

void insert(t_trie_node *root, char *word) {
    t_trie_node *node = root;
    int c;
    for (int i = 0; i < (int)strlen(word); i++) {
        c = (int)word[i];
        if (node->children[c] == NULL)
            node->children[c] = new_node();
        node = node->children[c];
    }
    node->is_end_of_word = true;
}

bool search(t_trie_node *root, char *word) {
    t_trie_node *node = root;
    int c;
    for (int i = 0; i < (int)strlen(word); i++) {
        c = (int)word[i];
        if (node->children[c] == NULL)
            return false;
        node = node->children[c];
    }
    return node->is_end_of_word;
}

int main() {
    t_trie_node *root = new_node();

    insert(root, "apple");
    insert(root, "apply");
    insert(root, "apple store");
    insert(root, "ape");
    insert(root, "banana");
    insert(root, "tomato");

    printf("%d\n", search(root,"apple")); // 1
    printf("%d\n", search(root,"app")); // 0
    printf("%d\n", search(root,"apex")); // 0
    printf("%d\n", search(root,"this is not inserted")); // 0

    return 0;
}

「単方向リストのchildrenが複数あって、各ノードで単語の終端フラグあるだけ」といえばそう。

*1:そのテストはかなり制約がゆるく、トライ木の代わりに全探索を書いたら通ってしまった

*2:各種エラーハンドリングは意図的に書いてない