けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 163 D - Sum of Large Numbers (緑色, 400 点)

AtCoder AtCoder400点 ABC-D 数え上げ問題 解空間:O(N^2)通りの選択肢 連結性に着目する 式変形 ダブルカウントを防ぐ場合分け ある量を固定して考える 等差数列 緑色diff N個からK個を選ぶ設定の問題 Greedy Greedy:各要素について独立に考えてよい 最大値や最小値に着目する 数学(代数) 考える整数が連続することに着目する

「作れる数が連続する整数になる」というの、実は結構よくある！！

問題へのリンク

問題概要

N+1 個の整数 10^{100}, 10^{100}+1, \dots, 10^{100}+N がある。

これらから K 個以上の整数を選んで合計して得られる整数としてありうるものの個数を 1000000007 で割ったあまりを求めよ。

制約

  • 1 \le N \le 2 \times 10^{5}
  • 1 \le K \le N+1

考えたこと

まず、

「この 10^{100} ってなんだよ！！！！！」

となるかもしれない。でも、よくよく考えると、これはそんなに恐れる必要はなくて...。 10^{100} だとでかすぎるので、試しに N = 5 として、 10^{100} の代わりに 10000 くらいにしてみよう。そうすると N+1 個の整数は

10000, 10001, 10002, 10003, 10004, 10005

になる。これらの整数から、いくつか選んで足してみる。このとき、

  • 1 個選んで足したときは、どのように選んでも、最上位の値が 1
  • 2 個選んで足したときは、どのように選んでも、最上位の値が 2
  • 3 個選んで足したときは、どのように選んでも、最上位の値が 3
  • 4 個選んで足したときは、どのように選んでも、最上位の値が 4
  • 5 個選んで足したときは、どのように選んでも、最上位の値が 5
  • 6 個選んで足したときは、どのように選んでも、最上位の値が 6

という感じになっている。たとえば 3 個の整数 10001, 10003, 10004 を選んで足すと、30008 になる。このように、最上位の値は 3 になるのだ。つまり、「異なる個数を選んで足した場合、それらが一致することは絶対にない」ということだ。よってこの問題は、

  • K 個選ぶ場合
  • K+1 個選ぶ場合
  • ...
  • N+1 個選ぶ場合

をそれぞれ独立に求めて、合計すれば OK！！！

k 個選ぶ場合を求める

よって問題は、 k 個選ぶ場合を求めることに帰着された。ここまで来ると、 10^{100} とか最早関係ない。さっきの例でいえば、

10001 + 10003 + 10004 = 30008

というのと、

1 + 3 + 4 = 8

というのは同じことなのだ。よって問題は、以下の問題に帰着される。

N+1 個の整数 0, 1, 2, \dots, N の中から k 個選んで足してできる整数は何通りあるか？

またしても、具体例を使って考えてみよう。 N = 5, k = 3 としてみる。0, 1, 2, 3, 4, 5 から 3 個選んでできる整数を並べてみよう。

  • 3 (0 + 1 + 2 など)
  • 4 (0 + 1 + 3 など)
  • 5 (0 + 1 + 4 など)
  • 6 (0 + 1 + 5 など)
  • 7 (0 + 2 + 5 など)
  • 8 (0 + 3 + 5 など)
  • 9 (0 + 4 + 5 など)
  • 10 (1 + 4 + 5 など)
  • 11 (2 + 4 + 5 など)
  • 12 (3 + 4 + 5 など)

出来上がる数は、このように 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 の 10 個になる。ここで注目したいことは、最小の数 3 から、最大の数 12 まで「連続する整数を作ることができる」ことだ。直感的には、これは一般に成り立ちそうだ。成り立つのだ (整数 a をいくつかの整数の和で表せるとき、それらのどれかを 1 増やせる感じ)。

よって、

  • k 個足してできる最小の整数を求める ( a とする)
  • k 個足してできる最大の整数を求める ( b とする)
  • 出来上がる整数は b - a + 1 個である

という風にできる。最小の整数は小さい順に k 個選べばよいので、

  • 0 + 1 + \dots + (k-1) = \frac{k(k-1)}{2}

となり、最大の整数は大きい順に k 個選べば良いので、

  • N + (N-1) + \dots + (N-k+1) = \frac{k(2N - k + 1)}{2}

となる。よって、求める個数は、

  • \frac{k(2N - k + 1)}{2} - \frac{k(k-1)}{2} + 1

となる。

まとめ

まとめると、 k = K, K+1, \dots, N+1 について、

  • \frac{k(2N - k + 1)}{2} - \frac{k(k-1)}{2} + 1

を合計すればよい。計算量は、各 k についての答えが O(1) で求められるので、 O(N)。

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1000000007;
int main() {
    long long N, K;
    cin >> N >> K;
    long long res = 0;
    for (long long k = K; k <= N+1; ++k) {
        long long first = k * (k-1) / 2;
        long long final = (N*2-k+1) * k / 2;
        long long add = final - first + 1;
        res = (res + add) % MOD;
    }
    cout << res << endl;
}

さらに O(1) へ

ここまでは必要ないけど、今回の問題は O(N) から O(1) に改善することもできる。まず、

\frac{k(2N - k + 1)}{2} - \frac{k(k-1)}{2} + 1 = -k^{2} + (N+1)k + 1

よって求める個数 A は、

A = \sum_{k = K}^{N+1} (-k^{2} + (N+1)k + 1)

となる。ここまで来れば、シグマ計算を頑張ればできそうだ。ちょっと見やすくするために、 M = N+1, L = K-1 とすると、

A = \sum_{k = L+1}^{M} (-k^{2} + Mk + 1) = \sum_{k = 1}^{M} (-k^{2} + Mk + 1) - \sum_{k = 1}^{L} (-k^{2} + Mk + 1)

と変形できる。見やすくなった。ここでシグマの公式を思い出そう。

  • \sum_{k = 1}^{M} k = \frac{M(M+1)}{2}
  • \sum_{k = 1}^{M} k^{2} = \frac{M(M+1)(2M+1)}{6}

これを使うと、

  • \sum_{k = 1}^{M} (-k^{2} + Mk + 1) = - \frac{M(M+1)(2M+1)}{6} + M\frac{M(M+1)}{2} + M = \frac{M(M^{2}+5)}{6}
  • \sum_{k = 1}^{L} (-k^{2} + Mk + 1) = - \frac{L(L+1)(2L+1)}{6} + M\frac{L(L+1)}{2} + L = \frac{-L(2L^{2} + 3L - 5 - 3M(L+1))}{6}

という風に計算できる。よって求める答えは、

A = \frac{M(M^{2}+5)}{6} + \frac{L(2L^{2} + 3L - 5 - 3M(L+1))}{6}

となる。

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1000000007;
int main() {
    long long N, K;
    cin >> N >> K;
    long long M = N+1, L = K-1;
    long long res = M*(M*M+5)/6 + L*(L*L*2+L*3-5-M*(L+1)*3)/6;
    cout << res%MOD << endl;
}