【Python】ABC403E Trie木をdataclassを使って表現してみる

はじめに

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を使う記述で問題ありませんでした。

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の問題を解く限りでは、特に問題はありませんでした。