AtCoder Beginner Contest 141のＣをPython3で解いたメモ

問題はこちら
C - Attack Survival

まずは愚直な方法で解く

N, K, Q = map(int, input().split(" "))
A = [int(input()) for i in range(Q)]

score = [K for i in range(N)]

for i in range(Q):
    num_ans_person = A[i] - 1
    score[num_ans_person] += 1
    for j in range(len(score)):
        score[j] += -1

for i in range(N):
    if score[i] > 0:
        print("Yes")
    else:
        print("No")

これだと計算時間がかかりすぎるためにTLE.
テストケースのsub1_13.txtはACだったが, 1285 msの時間がかかっている.

余事象側を考えてみる

正解数のほうに注目して高速化を図る.

N, K, Q = map(int, input().split(" "))
A = [int(input()) for i in range(Q)]

tmp = []
for i in range(N):
    tmp.append(A.count(i+1)) #正解数
# 得点は, 最初の点数 - 他人の正解数
# 他人の正解数は, 問題数 - 自分の正解数
# 得点は, 最初の点数 - 問題数 + 自分の正解数
# 得点 > 0 であれば, "Yes"
for i in range(N):
    if K - Q + tmp[i] > 0:
        print("Yes")
    else:
        print("No")

テストケースのsub1_13.txtでは, 264 msと, 高速になってはいるもののTLE.

本質的な部分で計算量が変わっていないのかと.
おそらくは, 下の部分の処理が素人くさくて遅い。

for i in range(N):
    tmp.append(A.count(i+1))

計算量を考えてみる

計算量を真面目に考える必要がありそう. O(N)みたいなやつ.

愚直な方法では, 問題数についてQ回の操作, 不正解の(N-1)人の点数を1点ずつ減点するので, O(QN).（多分）
N, Qがそれぞれ10^5なので, 最大10^10の計算量.

正解数に注目した方は, Q個のデータの入ったリストに対し, count関数でN回のループで計算.
これも同じような計算量になりそう.
forループを使わなかったことだけが高速化に効いてるだけっぽい.

毎回長さがQのリストを参照しているのが原因っぽいので直してみる

これでACでした.

N, K, Q = map(int, input().split(" "))    # N:人数, K:初期スコア, Q:問題数
A = [int(input()) for i in range(Q)]    # 正解者のリスト
B = [0 for i in range(N)]    # 各人の正解数

for i in range(Q):
    num_ans_person = A[i] - 1
    B[num_ans_person] += 1

# 得点は, 最初の点数 - 他人の正解数
# 他人の正解数は, 問題数 - 自分の正解数
# 得点は, 最初の点数 - 問題数 + 自分の正解数
# 得点 > 0 であれば, "Yes"
for i in range(N):
    if K - Q + B[i] > 0:
        print("Yes")
    else:
        print("No")