TrieのPHP実装書いてみた

はてなキーワードっぽいものの実装をする必要があるよねという話があって、前軽く書いたけど、
http://d.hatena.ne.jp/tsugehara/20120909/1347190811

tx-rubyの記事からたどるに、どーもこの手の処理を実装するにはTrieっていうアルゴリズムがいいらしい。
Google IMEにも使われているらしいぞ。

ここでTrieってアルゴリズムで書くのがいいらしいと。
PHPにはぱっと見組み込みの実装が無いので書く必要があるよねということで書いてみた。

ってこんなんでいいのか？なんかコストが想定よりでかいんだが。

<?php
class TTrie {
	var $trie = array();
	var $s;
	function __construct($s=NULL){
		$this->s = $s;
	}

	function add($s, $i=1, $l=NULL){
		if (! $i)
			throw new Exception('Invalid argument');
		if (! $l) {
			$l = strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (! array_key_exists($s, $this->trie)) {
				$this->trie[$s] = new TTrie($s);
			} else {
				$this->trie[$s]->s = $s;
			}
			return TRUE;
		}

		$c = substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			$this->trie[$c] = new TTrie();

		return $this->trie[$c]->add($s, $i+1, $l);
	}

	function search($s, $type='a') {
		$func = 'search'.strtoupper($type);
		return $this->$func($s);
	}

	function searchO($s, $i=1, $l=NULL){
		if (! $l) {
			$l = strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (array_key_exists($s, $this->trie))
				return $this->trie[$s];

			return FALSE;
		}

		$c = substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchO($s, $i+1, $l);
	}

	function searchA($s, $i=1, $l=NULL){
		if (! $l) {
			$l = strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (array_key_exists($s, $this->trie))
				return $this->trie[$s]->getAll();

			return FALSE;
		}

		$c = substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchA($s, $i+1, $l);
	}

	function getAll() {
		$ret = array();

		if ($this->s)
			$ret[] = $this->s;

		foreach ($this->trie as $k => $v)
			$ret = array_merge($ret, $v->getAll());

		return $ret;
	}
}
?>

基本的にはこんな感じで使う。

<?php
$trie = new TTrie();
$trie->add('こんにちは');
$trie->add('こんばんは');
$trie->add('おはよう');
$ret = $trie->search('こん');
?>

上記の例だと「こん」というキーワードを探して、「こんにちは」と「こんばんは」が取れる。
気になるのは、コスト的にUTF-8だとこんにちはで3*5の15バイト分もあるんだよね。再帰関数のコールが最低15回行われるという事。
searchには第二引数があって、もうちょっとコストの軽いこれで今回は実装をするつもりだけど、コール数自体は最低回数になるだけで最低回数より少なくなるわけではない。

<?php
$ret = $trie->search('こん', 'o');
?>

さらに厳密に1バイトずつ分解しているので、「こん」の段階でこれ以上探せなくなって、「こん」以下のデータ構造から「こんに」をたどって、「こんにち」をたどって、「こんにちは」をたどってようやく「こんにちは」の1ワード、さらに「こんば」をたどって「こんばん」をたどって「こんばんは」をたどってようやく「こんばんは」で2ワードとれる。
なんかあんまり早い気がしないんだけどどうなんだろう。どっか実装間違ってんのかね。

とりあえずWikiから得られる知識だとこんなもんが実装として妥当で、後は普通のsubstrなのかmb_substrなのかとか細部の最適化に余地はあるとしても、アルゴリズムとしては間違ってなさそうな気はすんだけど、どうなんだろね。
もうちょっと頑張れば「こん」の次が「こんにちは」と「こんばんは」で2階層にすることも出来るんだろうけど、あくまでダミーだしこんなもんでいいかなという気もする。

まあ、とりあえずこんなもんにして実装を進めてみる予定。

追記

ベンチマークをとってみると、mb系の方が少なくともUTF-8環境では早い事がわかったので、以後は基本的にこっちを使うことになると思う。

<?php
class TTrieMB {
	var $trie = array();
	var $s;
	function __construct($s=NULL){
		$this->s = $s;
	}

	function add($s, $i=1, $l=NULL){
		if (! $i)
			throw new Exception('Invalid argument');
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (! array_key_exists($s, $this->trie)) {
				$this->trie[$s] = new TTrieMB($s);
			} else {
				$this->trie[$s]->s = $s;
			}
			return TRUE;
		}

		$c = mb_substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			$this->trie[$c] = new TTrieMB();

		return $this->trie[$c]->add($s, $i+1, $l);
	}

	function search($s, $type='a') {
		$func = 'search'.strtoupper($type);
		return $this->$func($s);
	}

	function searchO($s, $i=1, $l=NULL){
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (array_key_exists($s, $this->trie))
				return $this->trie[$s];

			return FALSE;
		}

		$c = mb_substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchO($s, $i+1, $l);
	}

	function searchA($s, $i=1, $l=NULL){
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}
		if ($i == $l) {
			if (array_key_exists($s, $this->trie))
				return $this->trie[$s]->getAll();

			return FALSE;
		}

		$c = mb_substr($s, 0, $i);
		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchA($s, $i+1, $l);
	}

	function getAll() {
		$ret = array();

		if ($this->s)
			$ret[] = $this->s;

		foreach ($this->trie as $k => $v) {
			$ret = array_merge($ret, $v->getAll());
		}

		return $ret;
	}
}
?>

5万回の試行でローカルPCだとmb系が10秒くらい、mbじゃない方は14秒くらい。
まあスクリプト言語は組み込み関数を駆使した方が速度早いですよってだけで、アルゴリズム的に早い訳じゃないんだけどね。

追記2

メリットしかなかったので、キーを全文字ではなく一文字ずつにしてみた。

<?php
class TTrieMB {
	var $trie = array();
	var $s;
	function __construct($s=NULL){
		$this->s = $s;
	}

	function add($s, $i=0, $l=NULL){
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}

		$c = mb_substr($s, $i, 1);
		if ($i == ($l-1)) {
			if (! array_key_exists($c, $this->trie)) {
				$this->trie[$c] = new TTrieMB($s);
			} else {
				$this->trie[$c]->s = $s;
			}
			return TRUE;
		}

		if (! array_key_exists($c, $this->trie))
			$this->trie[$c] = new TTrieMB();

		return $this->trie[$c]->add($s, $i+1, $l);
	}

	function search($s, $type='a') {
		if (strlen($type) != 1)
			return FALSE;
		$func = 'search'.strtoupper($type);
		return $this->$func($s);
	}

	function searchO($s, $i=0, $l=NULL){
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}

		$c = mb_substr($s, $i, 1);
		if ($i == ($l-1)) {
			if (array_key_exists($c, $this->trie))
				return $this->trie[$c];

			return FALSE;
		}

		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchO($c, $i+1, $l);
	}

	function searchA($s, $i=0, $l=NULL){
		if (! $l) {
			$l = mb_strlen($s);
			if (! $l)
				throw new Exception('Invalid argument');
		}

		$c = mb_substr($s, $i, 1);
		if ($i == ($l-1)) {
			if (array_key_exists($c, $this->trie))
				return $this->trie[$c]->getAll();

			return FALSE;
		}

		if (! array_key_exists($c, $this->trie))
			return FALSE;

		return $this->trie[$c]->searchA($s, $i+1, $l);
	}

	function getAll() {
		$ret = array();

		if ($this->s)
			$ret[] = $this->s;

		foreach ($this->trie as $k => $v) {
			$ret = array_merge($ret, $v->getAll());
		}

		return $ret;
	}
}
?>

「こんにちは」を検索する時、「こ」→「こん」→「こんに」→「こんにち」→「こんにちは」ではなく、「こ」→「ん」→「に」→「ち」→「は」で探すようになってる。

追記3

GITに移します。
https://github.com/tsugehara/TrieMB