paizaの練習問題を解く総和の計算 Go編

こんにちは、ナナオです。

前回に引き続き競プロを実施していきたいと思います。

今回の問題は以下です。

総和の計算 | レベルアップ問題集 | プログラミング学習サイト【paizaラーニング】

実装

シンプルに実装するとこうです。

package main
import "fmt"
func main(){
    var a, b int
    fmt.Scan(&a, &b)
    r := 0
    for i := a; i <= b; i++ {
        r += i
    }
    fmt.Println(r)
}

ですがこれだと実行時間制約に間に合いません。

ということで解説を見ます。

1 - Nまでの総和はN(N + 1) / 2で求めることができるようです。

ということでAからBまでの総和は(1からBまでの総和) - (1からA - 1までの総和)ということになります。

この法則に従って実装を修正します。

package main
import "fmt"
func main(){
    var a, b uint64
    fmt.Scan(&a, &b)
    r := ((b * (b + 1)) / 2) - ((a * (a - 1)) / 2)
    fmt.Println(r)
}

と、修正しましたが、うまく値が出力されません。

どうやらこの問題の解答の最大値がGoだとオーバーフローしてしまうようです。

AIと壁打ちしたところ、math/bigパッケージを使えば解決するようです。

package main
import (
    "fmt"
    "math/big"
)
func main(){
    a := new(big.Int)
    b := new(big.Int)
    fmt.Scan(a, b)
    one := big.NewInt(1)
    two := big.NewInt(2)
    
    // b * (b + 1) / 2
    t1 := new(big.Int).Div(new(big.Int).Mul(b, new(big.Int).Add(b, one)), two)
    // a * (a - 1) / 2
    t2 := new(big.Int).Div(new(big.Int).Mul(a, new(big.Int).Sub(a, one)), two)
    r := new(big.Int).Sub(t1, t2)
    fmt.Println(r.String())
}

相当冗長なコードになりました。が、まぁ一応解決したのでいいでしょう。

ということで今回はこれで👍

実装#

実装