二次方程式の解を求める（桁落ちと精度に気を付けながら）

この記事はyukicoder「No.648　お　や　す　み」の解説です。
Python3で桁落ちと丸め誤差を避けるためのTipsです。

問題の概要

$i$秒目には、$i$匹の羊が柵を超えます。
柵を超えた羊の総数がちょうど$n$匹になるなら"YES"、ならないなら"NO"を出力してください。
$1\le n \le 2 \times 10^{18}$

シンプルに足していく手法

羊が柵を超えた総数を$i=1$から足していき、総数が$n$ちょうどになるか、$n$を超えるかを調べる方法

def simple(n, N):
  """ TLEします """
  sum_sheep = 0
  for i in range(1,N+1):
    sum_sheep += i
    if sum_sheep < n:
    continue
    elif sum_sheep == n:
      print("YES")
      print(i)
      exit()
    else:
      print("NO")
      exit()

if __name__ == "__main__":
  n = int(input())
  N = 2*(10**18)
  simple(n, N)

このような脳死アルゴリズムでは、 $n$ が非常に大きな値のとき、実行時間2秒以内の時間制限にかかり失敗します（TLEします）。

1からnまでの和の公式を使って求める手法

$1$から$i$までの和 $= i\times(i+1)/2$ を使って求めます。

def sum_one_to_n(n, N):
  """ 1からnまでの和 = n*(n+1)/2 を求める
  veryverylarge1.txtでTLE
  """
  for i in range(1, N+1):
    sum_sheep = i*(i+1)/2
    if sum_sheep < n:
      continue
    elif sum_sheep == n:
      print("YES")
      print(i)
      exit()
    else:
      print("NO")
      exit()
      
if __name__ == "__main__":
  n = int(input())
  N = 2*(10**18)
  sum_one_to_n(n, N)

これも最初の手法とやってることは変わりません。同じくTLEします。

二次方程式の解を求める手法（本題）

$1$から$i$までの和が$n$になるかを調べればよいわけですから、次の式が導出できます。
$i\times(i+1)/2 = n$
これを変形すると以下になります。
$i^{2}+i-2n=0$
つまり、上記の式を満たす$i$を探せば良いことになります（二次方程式の解を求める問題）。

import math

def _formula(n):
  """ 二次方程式の解を返す
  解が正になるのは問題文により自明です
  99_corner3.txtでWAになる（桁落ちが原因）
  """
  D = 1 + 8*n
  return (-1 + math.sqrt(D))/2  # ※式

def quadratic_formula(n):
  """ 二次方程式の解を求める方法
  1からnまでの和の公式を変形すると、
  i**2 + i - 2*n = 0
  このときのiを求めて、iが正の整数ならば"YES"である
  """
  i = _formula(n)
  if i != math.floor(i):
    print("NO")
    exit()
  print("YES")
  print(int(i))
  exit()

if __name__ == "__main__":
  n = int(input())
  N = 2*(10**18)
  quadratic_formula(n)

_formula()関数が二次方程式の解を求める関数です。

if i != math.floor(i):

の部分は、求めた解$i$が整数かどうかを判定しています。

この手法では、コーナーケース99_corner3.txtのテストで誤った解が求まります（WAになります）。原因は桁落ちです。
「二次方程式の解 桁落ち」でググればたくさんヒットしますが、計算機を使って二次方程式の解を求めるときには分子の有理化をすれば大抵の桁落ちを避けることが出来ます 。

二次方程式の解を求める（分子の有理化）

前回の手法の_formula()関数内の式を分子の有理化しました。それ以外は変えていません。

import math

def _formula(n):
  """ 二次方程式の解
  桁落ちに対処するため、_formula()の※式を分子の有理化したもの
  99_corner6.txtでWAになる
  """
  D = 1 + 8*n
  return 4*n/(1+math.sqrt(D))

この手法でも、残念ながらコーナーケース99_corner6.txtでWAになってしまいます（それ以外はACでした）。

有効桁数を大きくして計算精度を上げる

原因は_formula()関数内の除算の誤差です。

D = 1 + 8*n
return 4*n/(1+math.sqrt(D))

上記の計算結果で丸め誤差が生じています。

n = 1805000061750000527のとき、
4*n/(1+math.sqrt(D))

の計算結果は$1900000032.0$になります。整数なので、出力は"YES"です。

丸め誤差を避けるには、有効桁数を大きくして計算する必要があります。その場合はdecimalモジュールを使用します。

from decimal import *
getcontext().prec = 28  # 28桁はほとんどの問題解決に十分な大きさ
Decimal(4*n)/Decimal(1+math.sqrt(D))

計算結果は$1900000031.999999998947368439$になります。整数でないので、出力は"NO"です。