ABC163 C問題 with python3

ABC163 C問題 management

一言メモ

tleになったのでどこをどう変えればいいか考えてみる。tleになると時間制限さえなければ正解だったのに惜しいなと思ってしまいがちだが間違いは間違い。計算量を考えてアルゴリズムを組むことを意識しよう。

自分の誤答

N=int(input())
boss_list=list(map(int,input().split()))
num=0

for i in range(1,N+1):
    for j in range(i-1,N-1):
        if boss_list[j]==i:
            num+=1
    print(num)
    num=0

誤答のアルゴリズムは
社員番号２からNの全ての社員について上司の番号が１かどうか調べ、合致した回数を出力する.
社員番号3からNの全ての社員について上司の番号が2かどうか調べ、合致した回数を出力する.
...
社員番号N-1からNの全ての社員について上司の番号がN-2かどうか調べ、合致した回数を出力する.
社員番号NからNの全ての社員について上司の番号がN-1かどうか調べ、合致した回数を出力する.
という二重の繰り返し構造になっている。
この考え方だと無駄が多い。

約N回繰り返すfor文を二重に使っているのでオーダーが$O\left( N^{2}\right)$となりtleになった。
制約は$2\leqq N\leqq 2\times 10^{5}$なのでオーダーを$O\left( N\right)$に抑えたい。
出力はN回繰り返さなければいけないのでfor文でN回繰り返し１重で記述することを目指す。

ポイント

リストをうまく使うことで繰り返し処理が減らせる。
今回の場合では要素数N個のリストを作成し、x番目に社員番号x番の部下の数を格納することにすると以下のように完結にコードが書ける。

正解コード

N = int(input())
A = list(map(int, input().split()))

result = [0] * N
for a in A:
    result[a - 1] += 1
print('\n'.join(map(str, result)))

考察

今回リストを使うことで簡潔にプログラムが書けたが、どのような時に使うと有効なのだろうか？
複数の数値が入ったリストAに対してそれぞれの数値が要素として何個入ってるかをデータとして取り出したいときに新しく0がいっぱい入ったリストBを作り、リストBのx番目にAに数値xが何個入ってるかを格納していくといい。
わかりやすい例を以下に挙げる。
1から3までの数字が入ったリストA=[3,2,1,2,3,2,3,1]についてそれぞれの数が何個入ってるか調べたい、
このとき新しくリストB=[0]*3を用意して以下の処理を行うとBに求めたいデータが格納される。

A=[3,2,1,2,3,2,3,1]
B=[0]*3

for i in A:
    B[i-1]+=1
print(B)

#このときBは[2,3,3]となる。Aに含まれてる1,2,3の個数がそれぞれ求まった。

終わりに

今回の考察でわかったパターンは暗記しておくと他の問題を解くときにも使えそうだ。

複数の数値が入ったリストについてそれぞれの数値の要素の数を求める

ときには反射でリストを使えば書ける！と思いつきたい。