木構造のメモリ確保

トライ木を使って以下の問題にトライしたところ、TLEを食らいました。

http://poj.org/problem?id=3630

　最初に書いたコードはこれです。トライ木を構築して、語尾にあたる節点が子を持っていたらエラーとします。このエラー検出と同時に動的確保したメモリを解放しています。

struct Trie {
    bool tail;
    Trie *next[10];
    Trie() {
        tail = false;
        fill(next, next+10, (Trie*)0);
    }
};

Trie *find(char *t, Trie *r) {
    for (int i = 0; t[i]; ++i) {
        int c = t[i]-'0';
        if (!r->next[c]) r->next[c] = new Trie;
        r = r->next[c];
    }
    r->tail = true;
    return r;
}

bool validate(Trie *r) {
    bool hasCld = false;
    bool ret = true;
    for (int i = 0; i < 10; i++) {
        if (r->next[i]) {
            hasCld = true;
            ret &= validate(r->next[i]);
        }
    }

    if (r->tail && hasCld)
        ret = false;

    delete r;
    return ret;
}

int main() {
    int t, n;
    char num[12];

    scanf("%d", &t);
    while (t--) {
        scanf("%d", &n);
        Trie *rt = new Trie;
        REP(i, n) {
            scanf("%s", num);
            find(num, rt);
        }

        if (validate(rt))
            puts("YES");
        else
            puts("NO");
    }

    return 0;
}

　なんとか高速化する方法はないかと調べていたら、komiyamさんの日記におもしろいアイディアがあったので、参考にさせていただきました。毎回メモリを動的確保するのではなく、予め必要な分だけ静的領域に確保しておき、節点を生成するときは予め確保しておいたメモリのアドレスを渡すという方法です。これにより、メモリの確保/解放に係るオーバーヘッドがなくなります。また、メモリの解放をせずにすむために枝狩りができます。私の最初のソースでは、エラーを検出してもトライ木全体をトラバースしていました。これはメモリを解放するためです。メモリの解放がいらなくなると、エラーを検出した時点で探索をやめることができます。

struct Trie {
    bool tail;
    Trie *next[10];
    Trie() {
        tail = false;
        fill(next, next+10, (Trie*)0);
    }
};

Trie nodes[100000];
int ptr;

Trie *find(char *t, Trie *r) {
    for (int i = 0; t[i]; ++i) {
        int c = t[i]-'0';
        if (!r->next[c]) {
            nodes[++ptr] = Trie();
            r->next[c] = nodes + ptr;
        }
        r = r->next[c];
    }
    r->tail = true;
    return r;
}

bool validate(Trie *r) {
    for (int i = 0; i < 10; i++) {
        if (r->next[i]) {
            if (r->tail)
                return false;
            if (!validate(r->next[i]))
                return false;
        }
    }
    return true;
}

int main() {
    int t, n;
    char num[12];

    scanf("%d", &t);
    while (t--) {
        scanf("%d", &n);
        ptr = 0;

        nodes[ptr] = Trie();
        Trie *rt = &nodes[ptr];
        REP(i, n) {
            scanf("%s", num);
            find(num, rt);
        }

        if (validate(rt))
            puts("YES");
        else
            puts("NO");
    }

    return 0;
}

　以上がメイントピックですが、3点ほどちょっとしたTipsを書いておきます。

[C++のstructについて]

　C++のstructはCの構造体とは違います。C++の場合はstruct内にメンバ関数を持つことができます。classの場合アクセス修飾子を省略するとprivateになりますが、structの場合はpublicになるそうです。

http://www2s.biglobe.ne.jp/~ragnarok/program/cpp/struct.htm

[文字から数値への変換]

'0'　- '9'のcharを0 - 9のintに変換するとき、komiyamさんはおもしろいやり方をしています。

c & 15

でchar -> intに変換されます。確かに、'0'は16進数だと0x30なので、0x0fでマスクしてあげるとintに変換できます。

[ローカル変数のdelete]

最初、rootの接点だけローカル変数で宣言していて、validateでメモリ解放するとクラッシュしました。当然っちゃ当然ですがスタックに積まれた変数をdeleteしようとすると予期せぬ結果が発生します。再帰的処理で一括してメモリ解放するときなどに、ある特別なデータだけスタック変数だったとか結構ありそうな落とし穴なので注意。

http://stackoverflow.com/questions/441831/c-calling-delete-on-variable-allocated-on-the-stack