AtCorder|(復習)B - Counting of Trees

昨日、完答できなかった↓の問題の復習記事です。 atcoder.jp

自分の解答

n = int(input())
d = list(map(int,input().split()))
 
ans = 1
if min(d) != 0:
  ans = 0
d.sort()
for i in range(1,max(d)+1):
  ans *= d.count(i-1)**d.count(i)
  if ans == 0:
    break
print(ans%998244353)

条件を満たすN頂点の木の組み合わせの総数が (親の頂点の数**子の頂点の数)である規則性まで気づけたがTLE/WAとなり完答できませんでした。速度改善や例外対応が必要なのでAC者の模範解答を参考に上記のコードもリファクタしていきます。

ACしている人の解答例

from collections import Counter
mod=998244353
n=int(input())
d=list(map(int,input().split()))
q=Counter(d)
if q[0]!=1 or d[0]!=0:
    print(0)
    exit()
ans=1
for i in range(1,max(d)+1):
    ans*=pow(q[i-1],q[i],mod)
    ans%=mod
print(ans)

PythonistのAC者はcollectionsを使っているコードが多かったです。
AC者は、Counterクラスを使って入力dのリストから辺の数をkeyとした連想配列を最初に作っておくことで 、TLE対策していることがわかりました。

TLEの原因

#自分の解答
n = int(input())
d = list(map(int,input().split()))
 
ans = 1
if min(d) != 0:
  ans = 0
d.sort()
for i in range(1,max(d)+1):
  # ループのたびにcountを呼び出している
  ans *= d.count(i-1)**d.count(i)
  if ans == 0:
    break
print(ans%998244353)

count(ループのたびに辺の数に対する頂点の数を数え上げている)を呼び出していることがTLEの原因になっていました。

from collections import Counter
q=Counter(d)

を使って最初に、辺の数をkeyとした連想配列作る処理に修正してみます。

n = int(input())
d = list(map(int,input().split()))

from collections import Counter
q=Counter(d)

ans = 1
if min(d) != 0:
  ans = 0
d.sort()
for i in range(1,max(d)+1):
  ans *= q[i-1]**q[i]
  if ans == 0:
    break
print(ans%998244353)

一旦これで提出してみたが、4問WAが出てしまった。

WA(例外)の原因

例外ケースの考慮漏れが今回のWAの原因だった
以下の、場合では入力dの最小値が0以外の場合のみしか例外として0を出力できていない。

# 自分の例外抽出したコード
if min(d) != 0:

以下の２点を考慮すればWAとなる例外も対応できる
* 頂点からの辺が0となる が1つしか無い場合
* が0でない場合

最終的なリファクタコード

n = int(input())
d = list(map(int,input().split()))
 
from collections import Counter
q=Counter(d)

# 例外対応
if q[0]!=1 or d[0]!=0:
  print(0)
  exit()

ans = 1
for i in range(1,max(d)+1):
  ans *= q[i-1]**q[i]
  if ans == 0:
    break
print(ans%998244353)

解けそうで解けなかった問題は復習したほうがいいね。あー悔しい。