はじめに
ABC403のE問題でTrie木を使う問題があったので、dataclassを使ってTrie木を表現してみました。
AtCoderの問題を解くときは、タプルで表現して、力づくで解く等になりがち何ですが、dataclassを使うと、可読性が上がるので、一度試してみました。
Trie木とは
説明は省略します。
等が参考になると思います。
Trie木の実装
@dataclass
class TrieNode:
children: dict = field(default_factory=lambda: defaultdict(TrieNode))
is_x_terminal: bool = False
count_y: int = 0
ABC403 Eでは、文字列Xとして追加されたものかを判定するis_x_terminalと文字列Yの数をカウントするcount_yを持たせる必要があったので、この二つを持たせています。
この部分は課題に合わせて修正します。
使い方
例として、文字列Sを追加する処理を書いてみます
def add_string(trie: TrieNode, s: str) -> None:
node = trie
for c in s:
node = node.children[c]
node.is_x_terminal = True
trie_tree = TrieNode()
add_string(trie_tree, "abc")
add_string(trie_tree, "ab")
工夫した点
childrenの辞書は、子ノードがない場合には、TrieNodeで初期化する必要があります。
以下のコードのようにif文で記述する事になります。
def add_string(trie: TrieNode, s: str) -> None:
node = trie
for c in s:
if c not in node.children: # <-- ここが面倒
node.children[c] = TrieNode() # <-- ここが面倒
node = node.children[c]
node.is_x_terminal = True
これでも良いのですが、自動的に初期化するようにchildrenをdefaultdictを使って、初期化してあります。
子ノードがない場合でも、children[c]でアクセスするとその時点で、TrieNodeが作成されるので、if文でのチェックが不要です。
children: dict = field(default_factory=lambda: defaultdict(TrieNode))
TrieNodeの定義で、TrieNodeが使えている！
この自動的に初期化するコードは、TrieNodeの定義の中で、TrieNodeを使う必要があって、どうやって書けばいいのか悩みましたが、lamba式を使うことで、解決しました。
lambda式は、実行した時に初めて内容が評価されるため、lambda式の中で、TrieNodeを使う記述で問題ありませんでした。
Nodeのサーチ
defaultdictを使うことで、子ノードの自動追加するようになっているので、
ノードをサーチする際には、意図せず、子ノード作成されてしまわないように、ノードが存在しない事を確認して、存在しない場合は、サーチを停止するようにしておきます。
このif文は、省略しようがないと思います。
def search_string(trie: TrieNode, s: str) -> bool:
node = trie
for c in s:
if c not in node.children:
return False
node = node.children[c]
return True
TrieNodeを使った回答例
ABC403 E
from dataclasses import dataclass, field
import sys
from collections import defaultdict
input = lambda: sys.stdin.readline().rstrip("\r\n")
I = input
II = lambda: int(I())
LI = lambda: list(input().split())
LII = lambda: list(map(int, input().split()))
@dataclass
class TrieNode:
children: dict = field(default_factory=lambda: defaultdict(TrieNode))
is_x_terminal: bool = False
count_y: int = 0
trie_tree = TrieNode()
def add_x(S):
node = trie_tree
for s in S:
node = node.children[s]
if node.is_x_terminal:
# 既に追加済みの場合は処理不要
return 0
node.is_x_terminal = True
ret = node.count_y
add_count_y(S,-ret)
return ret
def add_y(S):
node = trie_tree
for s in S:
node = node.children[s]
if node.is_x_terminal:
# Prefixが一致するXが既にある場合追加しない
return 0
add_count_y(S, 1)
return 1
def add_count_y(S,val):
"""Trie木の根からSの間の全てのノードにcount_yにvalを足す"""
node = trie_tree
for s in S:
node = node.children[s]
node.count_y += val
Q =II()
ans = 0
for _ in range(Q):
t,s = LI()
if t == "1":
ans -= add_x(s)
else:
ans += add_y(s)
print(ans)
まとめ
- Trie木をdataclassを使って表現してみました。
- Trie木の、Nodeにもたせる情報をdataclassを使うことで、わかりやすく記述できました。
- 実行速度については、未評価です。ABC403 Eの問題を解く限りでは、特に問題はありませんでした。