競技プログラミング TLEについて

Question

前提・実現したいこと

AtCoder 041 C
https://atcoder.jp/contests/abc041/tasks/abc041_c

問題を解き、提出するとTLEとなってしまいました。
自分のコードのどこが遅いのか分からず修正に手間取っています。

分かる方、ご教授願いたいです。

発生している問題・エラーメッセージ

TLE

該当のソースコード

n = int(input())
a = list(map(int,input().split()))
A = list(reversed(sorted(a)))
 
for i in range(n):
  print(a.index(A[i])+1)

Answer 1

問題によると、生徒数Nは最大で10の5乗分です。高い方順に並べ替えた上でindexメソッドで探索した場合、最大でO(N^2)の計算量になるからではないでしょうか。つまり、10の10乗。Pythonですと（また、恐らくC++でも）この計算量だとTLEでしょう。

teratailでも過去、TLEに関して質問回答が挙がっていますので、参考になると思います。

競技プログラミングの解答がTLEになってしまう - teratail#121345
競プロの回答がTLEとなる - teratail#205829
AtcoderでTLEとなる原因について【ABC問題】 - teratail#212668

---

** 追記しました：2020-05-16 01:40 **

私自身の回答が気になったので検証してみました。質問者さんのオリジナルのコードをもとに、indexメソッドを使わない方法で試してみます。

まず、AtCoderの各コンテストのテストケースが収められているAtCoder のテストケースをあたってみると、残念ながらABC041のテストケースは無いようでした。その為、テストケースを作成するコードを作りました。簡単にするために、身長は出席番号の取り得る値と同じ(10の5乗)としています。

import random

n = pow(10, 5)
l = [str(v) for v in range(1, n + 1)]

print(n)
random.shuffle(l)
print(' '.join(l))

このプログラムを実行すると、以下のように生徒100000人分の身長ダミーデータを出力します。

100000
80228 51819 9780 60213 55267 31769 29648 4721...(100000件のデータ)

このデータ（ファイル in1.txt）をもとに、まず質問者さんオリジナルのコードで実行し、実行時間をtimeコマンドで計測してみます。尚、実行環境としてはWindows10/WSL(Ubuntu) 、CPUはIntel Core i7-4870HQ 2.5GHz です。

$ time python3 x1.py < in1.txt > out1.txt 

real    0m59.463s
user    0m59.422s
sys     0m0.031s

およそ59.5秒かかります。AtCoderでは完全にTLEです。

このコードをもとにindexメソッドを使用せず、身長をキーとして、出席番号(-1)を値とした辞書（hash）で検索するよう実装してみます。

n = int(input())
a = list(map(int,input().split()))
A = list(reversed(sorted(a)))

kv = {}
for i in range(n):
  kv[a[i]] = i

for i in range(n):
  print(kv[A[i]] + 1)

これを同じデータでテストします。

$ time python3 x1b.py < in1.txt > out1b.txt 
real    0m0.208s
user    0m0.141s
sys     0m0.031s

実行時間が0.2秒強、となりました。AtCoderのサイトで実際に提出した訳ではありませんが、これであればACできそうです。それぞれ、出力結果も同じです。

$ ls -l out1.txt out1b.txt 
-rwxrwxrwx 1 user01 user01 588895 May 16 01:17 out1.txt
-rwxrwxrwx 1 user01 user01 588895 May 16 01:23 out1b.txt
$ diff  out1.txt out1b.txt 
$ 

参考になれば幸いです。

Comments

[labpynguin]

詳しくありがとうございます。

ようやく理解できました。
ありがとうございました。