Edit

【pythonでABC190を解説】D - Staircase Sequences

問題概要

問題ページ

D - Staircase Sequences
D - Staircase Sequences

問題ページへ移動する

問題文

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

制約

  • \(1 ≤ N ≤ 10^{12}\)
  • \(N\) は整数

問題の考察

貪欲に計算すると\(1 ≤ N ≤ 10^{12}\)なので間に合わない。

等差数列なので、とりあえず等差数列の和の公式が思い浮かぶ。

等差数列の和の公式

  • \(\displaystyle　\frac{項数(初項+末項)}{2} = 等差数列の和\)
  • \(末項=初項+項数-1\)
  • \(項数=n,初項=a,末項=b,等差数列の和=N\)とする

この式を変形する。

式の変形

  • \(\displaystyle \frac{n(a+a+n-1)}{2}=N\)
  • \(n(2a+n-1)=2N\)
  • \(\displaystyle 2a+n-1=\frac{2N}{n}\)
  • \(\displaystyle a=\frac{\frac{2N}{n}-n+1}{2}\)

ここで\(a\)は整数なので、if (2 * n) % i == 0:とif ((2 * n) / i - i + 1) % 2 == 0:として\(a\)が整数か判別している。

（if ((2 * n) / i - i + 1) % 2 == 0:だけで良いじゃん！というツッコみと変数が説明と違うじゃんというツッコみはしないでください・・・）

\(a\)が整数の場合に該当の数列が存在することが分かる。

しかし、これでは\(1 ≤ N ≤ 10^{12}\)なので間に合わないので工夫する。

計算量の削減

  • \(\displaystyle 2a-1= \frac{2N}{n}-n\)と式を変形
  • \(\displaystyle \frac{2N}{n}-n\)の部分を考える
  • 入力例1の場合
    • \(n=1\)の時\(23\)
    • \(n=3\)の時\(5\)
    • \(n=8\)の時\(-5\)
    • \(n=24\)の時\(-23\)
  • \(\displaystyle \frac{2N}{n}\)で\(2N\)を割った場合、符号が逆になる
    • 例えば、\(n=3\)と\(\displaystyle n=\frac{24}{3}=8\)の商は符号が逆
    • \(\displaystyle \frac{2N}{n}-n\)を考えれば明らか
  • 符号が逆
    • 偶奇が等しい
    • 偶奇が等しい値に\(-1\)したものの偶奇も等しい
    • 片方が割り切れるなら、もう片方も割り切れる
  • \( \sqrt{2N}\)まで計算すれば良い
    • \(n\)が条件を満たせば、\(\displaystyle \frac{2N}{n}\)も条件を満たす

というような感じで結局は\( \sqrt{2N}\)まで計算すれば条件を満たす数列の個数はカウントできる。

これならば\( \sqrt{10^{12}}=10^6\)なので間に合う。

たびすけ
たびすけ
等差数列の式の変形と、計算量の削減方法の\(2\)つを考える問題。
入力例をよく見ると計算量の削減方法の仮説が立てられますが、コンテスト中に証明しようとするとちょっと時間がかかるかもしれません・・・

ACコード

import sys


def solve():
    input = sys.stdin.readline
    mod = 10 ** 9 + 7
    n = int(input().rstrip('\n'))
    cnt = 0
    for i in range(1, int((n * 2) ** 0.5) + 1):
        if (2 * n) % i == 0:
            if ((2 * n) / i - i + 1) % 2 == 0:
                cnt += 2
    print(cnt)


if __name__ == '__main__':
    solve()

-Edit
-