O(1)のはずの関数呼び出しでTLE

Question

#問題

ABC159 D問題より

#時間切れコード

from itertools import groupby
from math import factorial
def comb(n, k):
    if n >= k:
        return factorial(n) // (factorial(k) * factorial(n - k))
    else:
        return 0

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

dic = {}
for i in range(1, N + 1):
	dic[i] = 0
for a in A:
	dic[a] += 1

v_sum = 0
for k, v in dic.items():
	v_sum += comb(v, 2)

for a in A:
	ans = v_sum - comb(dic[a], 2) + comb(dic[a] - 1, 2)
	print(ans)

#正解コード

from itertools import groupby
from math import factorial
def comb(n, k):
    if n >= k:
        return factorial(n) // (factorial(k) * factorial(n - k))
    else:
        return 0

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

dic = {}
for k, v in groupby(sorted(A)):
	dic[k] = len(list(v))

v_sum = 0
for k, v in dic.items():
	v_sum += comb(v, 2)

for a in A:
	ans = v_sum + 1 - dic[a]
	print(ans)

#質問
コードを書き換えたのは下から2行目の部分のみです。comb(dic[a], 2) + comb(dic[a] - 1, 2)だと見栄えが悪いなと思い手で計算をして1 - dic[a]としたら思いもよらず実行時間が縮まった次第です。ループを回すような関数であればともかく、今回用いたcomb関数は計算量O(1)でおさまるからfor a in Aのループに入れておいても大丈夫だろうと高を括っていたところ、実行時間は2秒を超過してしまいましたが、dicへのアクセスのみに変えることで実行時間が1.1秒ほどになりました。forループの中でcombを呼び出しただけでどうしてそれほどの計算量負荷がかかっていたのでしょうか。素人質問にて恐縮ですが、ご教授のほどお願い申し上げます。

Answer 1

###提示ソースとは別物ですが
（２秒以内で処理できるかはわかりませんが）私ならば全体をソートの上隣り合うものを比較することをまず考えます。

Answer 2

今回用いたcomb関数は計算量O(1)でおさまるから

factorialの計算量次第です。

Comments

[maisumakun]

多倍長整数となると、四則演算自体もO(1)とはみなせなくなります。

[kay_ventris4]

定数倍のサイズについても考慮しなければいけなかったのですね。ありがとうございました。

[episteme]

「いかなる N に対しても 常に8時間かかる」でも O(1) やしね。

[hayataka2049]

https://github.com/python/cpython/blob/main/Modules/mathmodule.c#L1852
こんな実装です。計算量の評価は私にはできませんが、とりあえずコメントを読めばアルゴリズムの概略がわかります。

ついでに速度をipythonで実測。
from math import factorial

%%timeit
factorial(10**2)
1.28 µs ± 399 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%%timeit
factorial(10**3)
69.6 µs ± 682 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
factorial(10**4)
4.52 ms ± 1.19 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
factorial(10**5)
198 ms ± 66.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

1桁増やすとラフに50倍前後の時間が持っていかれる感じ。

[kay_ventris4]

お二方ともご丁寧にありがとうございます。非常に勉強になりました。