ABC437Eを解いた【trie木】

筆者はレート800前後の茶～緑コーダ

ABC437当日にACできなかったE問題を解いていく

実装コード

Trie木という構造を用いて、
すべての配列を木構造の中で追えるようにする

同じ配列が同じノードになるように添字を管理しつつ、
辞書順になるように工夫しながらDFSを行う

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)
                
def main():
    N = rI()   
   
    trie = [defaultdict(lambda:-1) for _ in range(N+1)]
    
    K = [[] for _ in range(N+1)]
    
    P = defaultdict(list)
    T = [-1]*(N+1)
    T[0] = 0
    j = 1
    for i in range(N):
        x, y = rLI()
        tx=T[x]
        if trie[tx][y] == -1:
            trie[tx][y] = j
            j += 1
            K[tx].append(y)
        t = trie[tx][y]
        T[i+1] = t
        P[t].append(i+1)
            
    ans = []
    stack = deque([0])
    while stack:
        cur = stack.pop()
        ans += P[cur]
        for nxt in sorted(K[cur],reverse=True):
            stack.append(trie[cur][nxt])
    print(*ans)
                
    
if __name__ == '__main__':
    main()

感想

コンテスト時間内で似たような概念を作ろうとしてうまく行かなかったが、
解説を読んで新たなデータ構造を学ぶことができた。