Go で ARC091 D - Remainder Reminder を解いてみました。

問題

問題文

高橋君は、N 以下の正の整数の 2 つ組 (a,b) を持っていましたが、忘れてしまいました。
高橋君は、a を b で割ったあまりが K 以上であったことを覚えています。高橋君が持っていた組としてあるうるものの個数を求めてください。

制約

$ 1 \le N \le 10^5 $
$ 0 \le K \le N-1 $
入力はすべて整数である

( https://atcoder.jp/contests/arc091/tasks/arc091_b より引用)

考え方

$ 1 \le a,b \le N $ の範囲で全列挙したりして考察しましたが、全然思いつかずギブアップ。数学系は苦手です。 $a$ を固定して考えたのが良くなかった模様。

公式解説によれば、ある $b$ に対して $a$ を $b$ で割った余り、つまり $a % b$ が $K$ 以上である、 $1 \le a \le N$ の範囲の $a$ の個数を $O(1)$ で求めることが可能です。

b を固定して考えましょう。b = 1, 2, · · · N に対し、a を b で割った余りが K 以上であるような 1 ≤ a ≤ Nの個数が高速に求められれば良いです。簡単のため、a = 0 も許すことにして、あとで a = 0 の場合 (これは簡単に求められます) を引くことにしましょう。
さて、整数 p, q を用いて N = pb + r(0 ≤ r < b) という形で N を一意的に表したとき、a を 0 から N まで順に動かせば、a を b で割った余りを順に並べたものは、0, 1, 2, · · · b − 1 という列が p 回繰り返され、最後に 0, 1, 2, · · · r という列が付け加わったものになります。
0, 1, 2, · · · b − 1 という列が p 回繰り返される部分には条件を満たす a の個数は p × max(0, b − K) 個、最後の部分には max(0, r − K + 1) 個あるので、条件を満たす a の個数が O(1) 時間で求められ、O(N) 時間でこの問題を解くことができました。
( https://img.atcoder.jp/arc091/editorial.pdf より引用)

一度読んだだけではよくわからなかったので、具体例で考えてみます。

具体例

入力を $N = 100, K = 4$ とします。$b$ に適当な数をあてて考えてみます。
$b = 7$ として $a$ を $0$ から $100$ まで可変させたときの $r = a%b$ と、 $b = 12$ として同様の $r$ の値を次の図1に示します。

$b=7$ の場合

まず $b=7$ の場合、 $r$ は $0,1,2,3,4,5,6,0,1,2,3,4,...$ と $0$ から $6$ までの数(図1において、緑(色がわかりにくいですが)の波括弧でくくった部分)を繰り返すことになります。

ひとつの繰り返しの中を見てみると、このなかで $r$ が $K=4$ 以上であるものは、 $4,5,6$ の3つ、つまり $b-K$ 個です（図1の赤い四角で囲んだ部分）。ただし、そもそも $b$ が $K=4$ 以下、つまり $b=3$ のような場合に0でり、負の数になってしまったら困るので、正確には $max(0, b-K)$ 個になります。

この繰り返しの部分の個数は明らかに $N/b = 14$ 個です。つまり $b=7$ において、この繰り返し全体で条件を満たす $r$ の数は $N/b \times max(0, b-K) = 14 \times 3 = 42$ となります。

最後の $a=98, 99, 100$ のときに $r$ には $K=4$ 以上のものはないので、条件を満たす $r$ の数は $0$ 個です。

つまり、 $b=7$ の場合に条件を満たす $a$ の個数は $42+0 = 42$ 個になります。

b=12 の場合

次に、 $b=12$ の場合、を考えます。
先程と同じように $r$ が $0,1,2,3,4,5,6,7,8,9,10,11$ のパターンを繰り返すことになり、条件を満たす $r$ の数は $N/b \times max(0, b-K) = 8 \times 8 = 64$ 個です。

最後の $a=96, 97, 98, 99, 100$ のとき $r$ には $K=4$ 以上のものがありますが、これは $N$ を $b$ で割った余りの個数 $N%b$ と、その直前の余り0の1個の中から、$K=4$ 以下の個数を引き算して考えれば良いので、 $N%b+1-K$ で計算できます。ただし、$N%b+1$ が $K=4$ 未満の場合に負にならないよう、max(0, N%b+1-K) とすれば正確にカウントできます。ここでは $max(0, N%b+1-K) = max(0, 4+1-4) = 1$ 個になります。

つまり、 $b=12$ の場合に条件を満たす $a$ の個数は $42+0 = 42$ 個になります。

このようにして、 $b$ を $1$ から $N$ まで可変させて、それぞれの $b$ に対する $a$ の個数をカウントすることで、答えが求まります。

注意点

この方法では $a=0$ の場合も含めて考えることになります。もし、 $a=0$ （どの数で割っても余りが $0$ ）が含まれるような $K$, つまり $K=0$ の場合には $a=0$ の分を減じてやる必要があります。これは図1を見ると明らかなとおり、すべての b に一つずつありますから、全体としては $N$ 個になります。

よって、$K=0$ のときのみ答えから $N$ を引いてやればOK.

$K=0$ の場合、すべての $(a,b)$ の組が条件を満たすので、答えである $N^2$ を出力して終了、でも良いです。

提出コード

最終的に次のようなコードを書きました（抜粋）。全体はこちら

func main() {
	defer stdout.Flush()
	N := readInt64()
	K := readInt64()

	var ans int64
	for b := int64(1); b <= N; b++ {
		ans += N / b * max(0, b-K)
		ans += max(0, N%b+1-K)
	}

	if K == 0 {
		ans -= N
	}

	println(ans)
}

func max(a, b int64) int64 {
	if a < b {
		return b
	}
	return a
}