paizaラーニングレベルアップ問題集の「文字列収集」をやってみた

paiza×Qiita記事投稿キャンペーンということで、paizaラーニングレベルアップ問題集の文字列収集をやってみました。

---

入力の後半でクエリ文字列が来てから市場に出回っている$N$個の文字列を精査していては、時間がかかりすぎてしまいます。予め、このクエリ文字列が来たら必要金額はいくら、ということを計算しておきましょう。クエリ文字列と必要金額は、連想配列に格納しておきます。（下の表の、太字の入力に対して、細字の文字列と数値を格納。同じ文字列（キー）が来たら数値を足し合わせる）

bcac	3
b	3
bc	3
bca	3
bcac	3

#include <iostream>
#include <map>
#include <utility>
using namespace std;

int main() {
	int n, m;
	cin >> n >> m;
	map<string, int> dict;
	for (int i = 0; i < n; i++) {
		string s;
		int p;
		cin >> s >> p;
		for (int j = 1; j <= s.size(); j++) {
			dict[s.substr(0, j)] += p; //s.substr(0,j)がdictに登録されていない場合はdict[s.substr(0,j)]に0+pが格納される。
		}
	}
	for (int i = 0; i < m; i++) {
		string q;
		cin >> q;
		cout << dict[q] << endl;
	}
	return 0;
}

C++のmapは、連想配列に登録されていないキーに対する値を求めようとしたとき、デフォルト値が返されます（値がint型の場合は0）。Pythonのdefaultdictも同様です。しかし、多くの言語では、連想配列に登録されていないキーに対する値を求めようとするとエラーになりますので、既にキーが登録されている/されていないでチェック＆分岐処理が必要になります。

---

以上で終わりなのですが、トライ木というものを使って実装してみましょう。よく分からないという方は、リンク先やタグ遷移先を参照してください。
トライ木メニューを参考にしてトライ木を実装してみます。今回は、この問題を解けるように改変しています。

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

typedef struct Node {
	int value; //合計必要金額
	struct Node *child[26];
} Node;

typedef struct Trie {
	Node *root;
} Trie;

void insert(Trie *trie, const char *key, int value) {
	Node *node = trie->root;
	for (const char *c = key; *c; c++) {
		int ch = (int) (*c - 'a');
		if (!node->child[ch]) {
			if (!(node->child[ch] = (Node*) malloc(sizeof(Node)))) {
				abort();
			}
			memset(node->child[ch], 0, sizeof(Node));
		}
		node->child[ch]->value += value;
		node = node->child[ch];
	}
}

int search(const Trie *trie, const char *key) {
	Node *node = trie->root;
	for (const char *c = key; *c; c++) {
		int ch = (int) (*c - 'a');
		if (!node->child[ch]) {
			return 0;
		}
		node = node->child[ch];
	}
	return node->value;
}

void _clear(Node *node) {
	for (char c = 'a'; c <= 'z'; c++) {
		if (node->child[c - 'a']) {
			_clear(node->child[c - 'a']);
		}
	}
	free(node);
}

void clear(Trie *trie) {
	_clear(trie->root);
	trie->root = NULL;
}

---

上のトライ木を用いて、問題を解いていきます。

int main() {
	int n, m;
	scanf("%d %d", &n, &m);

	Trie trie;
	trie.root = (Node*) malloc(sizeof(Node));
	memset(trie.root, 0, sizeof(Node));
	for (int i = 0; i < n; i++) {
		char s[101];
		int p;
		scanf("%s %d", s, &p);
		insert(&trie, s, p);
	}
	for (int i = 0; i < m; i++) {
		char q[101];
		scanf("%s", q);
		printf("%d\n", search(&trie, q));
	}
	clear(&trie);
	return 0;
}

---

問題文の入力例1では、まず、以下のようにトライ木に格納されます。
bcac 3

abcd 14

abccjg 92

bcaddgie 2

abcd 6

cd 200

---

そして、以下のように値を取り出していきます。丸ノードの値が答えです。
b

---

a

---

abcd

---

gagioheo

gagioheoで始まる単語は格納されていませんので、0が答えです。

---

cb