AtCoder-ABC190 D - Staircase Sequences【Python解答例】

コンテストでは解けなかったAtCoder Beginner Contest190のD問題を記事にしていきます。
AtCoder Beginner Contest 190 - AtCoder

AtCoder Beginner Contest190 D - Staircase Sequences

D - Staircase Sequences

問題文

整数からなる公差 1 の等差数列のうち、総和が N であるものはいくつあるでしょうか？

制約

・1≤N≤10^12
・N は整数

解答例

import math

n = int(input())

def check(z):
    tmp = n2 // z - z + 1
    if tmp % 2 == 0:
        return True

ans = 0
n2 = 2 * n
N = int(math.sqrt(n2))
for x in range(1,N+1):
    if n2 % x == 0:
         y = n2 // x
         if check(x):
             ans += 1
         if check(y):
             ans += 1
else:
    print(ans)

解説

問題はシンプルに問題文通り総和がNになる公差1の等差数列の数を答える問題です。
公式解説動画の解法になります。

動画に紹介されていましたが、公差1の等差数列の和を式変形していくと初項が整数になるための2つの条件が出てきます。
ひとつが2N/lが整数であること（Nは等差数列の和、lは数列の長さ）
もうひとつは{(2N/l) - l + 1}が偶数であることです。

2N/lが整数である必要があるので、ループは√2Nまで調べればよいようです。
これで計算量が大きく減りますのでTLEを回避できます。

ループでは2N/lがxで割り切れるか調べて、割り切れるなら もうひとつの約数であるyを求めます。
あとはxとyそれぞれで{(2N/l) - l + 1}が偶数であるか調べ、偶数ならans変数に1を足します。

ループが抜けた後にans変数を出力すればOKでした。

解けはしませんでしたが個人的にD問題はおもしろかったです。
式変形しながらどうすれば解けるか考えるのが楽しいと思えました。
こんな問題もたまに出てくれたらいいですね。

また、約数の場合は計算量が√になることは勉強になりました。
今後に活かしていきたいです。

次こそはD問題をコンテストで解きたい…

ABC190の関連記事はこちら
ebisuke33.hatenablog.com
ebisuke33.hatenablog.com