AtCoder Beginner Contest 163 C を Python で

AtCoder Beginner Contest 163 C を Python で解説していきたいと思います。
それでは，さっそくまずは問題文から。
https://atcoder.jp/contests/abc163/tasks/abc163_c

C – management

実行時間制限: 2 sec / メモリ制限: 1024 MB

配点 : 300 点

問題文

N 人の社員からなる会社があり、各社員には 1,…,N の社員番号が割り当てられています。

社員番号 1 の社員以外の全ての社員には、自分より社員番号が小さい直属の上司がちょうど 1 人います。

X さんが Y さんの直属の上司であるとき、Y さんは X さんの直属の部下であるといいます。

社員番号 i の社員の直属の上司の社員番号が Ai であることが与えられます。各社員について直属の部下が何人いるか求めてください。

制約

2≤N≤2×10^5
1≤Ai<i

入力

入力は以下の形式で標準入力から与えられる。

N
A2 … AN

出力

社員番号 1,2,…,N のそれぞれの社員について、直属の部下が何人いるか、改行区切りで出力せよ。

入力例 1

5 1
1 2 2

出力例 1

2
2
0
0
0

社員番号 1 の社員の直属の部下は社員番号 2,3 の 2 人です。

社員番号 2 の社員の直属の部下は社員番号 4,5 の 2 人です。

社員番号 3,4,5 の社員には直属の部下はいません。

入力例 2

10
1 1 1 1 1 1 1 1 1

出力例 2

9
0
0
0
0
0
0
0
0
0

入力例 3

7
1 2 3 4 5 6

出力例 3

1
1
1
1
1
1
0

解説

少し問題文の意図が読み取りずらい問題ですが，結局は与えられる A2~AN の値の中に，１～Nの値はそれぞれいくつあるかを出力するという問題です。
まず一番簡単に思いつく方法は次のようなコードかと思います。

n = int(input())
a = list(map(int, input().split()))
 
for i in range(1,n+1):
    print(a.count(i))

つまり，そのまま１～Nまでの値 i について，A のリストの中にその数 i は何個あるかを count 関数を使って出力するというものです。
しかしこれだと O(n^2) となり TLE になってしまいます。
そこで，次のような実装を考えます。

まず，あるリスト b を作ります。これは０が n 個あるリストで，A の中の1~N までの値 i の数を表します。
これに for 文を用いて A の中の全ての数を調べ，その数 i に対応する b[i-1] に１を加えます。
あとは b の値をそれぞれ出力してやればよいです。
これだと O(n) でこの問題が解けました。
それでは，以下に解答例を載せておきます。

解答例

n = int(input())
a = list(map(int, input().split()))
 
b = [0]*n
 
for i in a:
    b[i-1] += 1
 
for i in b:
    print(i)

いかがだったでしょうか。
O(n^2) だと TLE になるので，計算量を減らすことを考えることが肝になりました。
何かありましたらコメントしてくださいね！
他の問題も解説しているので，見ていってください！
それでは！

タムログ

プログラミングについての情報発信をしていきます！