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

# はじめに

ABC403のE問題でTrie木を使う問題があったので、dataclassを使ってTrie木を表現してみました。

AtCoderの問題を解くときは、タプルで表現して、力づくで解く等になりがち何ですが、dataclassを使うと、可読性が上がるので、一度試してみました。

# Trie木とは

説明は省略します。

https://algo-logic.info/trie-tree/

等が参考になると思います。

# Trie木の実装

```python
@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を追加する処理を書いてみます

```python
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文で記述する事になります。

```python
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文でのチェックが不要です。

```python
children: dict = field(default_factory=lambda: defaultdict(TrieNode))
```

# TrieNodeの定義で、TrieNodeが使えている！

この自動的に初期化するコードは、TrieNodeの定義の中で、TrieNodeを使う必要があって、どうやって書けばいいのか悩みましたが、lamba式を使うことで、解決しました。
lambda式は、実行した時に初めて内容が評価されるため、lambda式の中で、`TrieNode`を使う記述で問題ありませんでした。

# Nodeのサーチ

`defaultdict`を使うことで、子ノードの自動追加するようになっているので、
ノードをサーチする際には、意図せず、子ノード作成されてしまわないように、ノードが存在しない事を確認して、存在しない場合は、サーチを停止するようにしておきます。
このif文は、省略しようがないと思います。

```python

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を使った回答例

<details><summary>ABC403 E</summary>

```python
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)
```
</details>

# まとめ

- Trie木をdataclassを使って表現してみました。
- Trie木の、Nodeにもたせる情報をdataclassを使うことで、わかりやすく記述できました。
- 実行速度については、未評価です。ABC403 Eの問題を解く限りでは、特に問題はありませんでした。