AtCoder精選良問「D - Sum of Large Numbers」Diff 960

2023年5月3日 00:27

この記事はAtCoderの精選良問集マガジンの中の記事です。
私が良問だと思った問題を解説付きでまとめています。

問題リンク

問題概要

10^100, 10^100 + 1, ..., 10^100 + N からなるN + 1 個の数列が与えられる。この数列の中から K 個以上の数を選ぶとき、その和としてあり得るものの個数を mod(10^9 + 7) で求める、という問題。

考え方

要するにN+1個の数列からいくつか選んで足したときに作りうる数はいくつある？ということを聞かれている。

まずは入力例１の解説を見てみる。

入力はN=3でK=2

(10^100) + (10^100 + 1) = 2 × 10^100 + 1
(10^100) + (10^100 + 2) = 2 × 10^100 + 2
(10^100) + (10^100 + 3) = (10^100 + 1) + (10^100 + 2) = 2 × 10^100 + 3
(10^100 + 1) + (10^100 + 3) = 2 × 10^100 + 4
(10^100 + 2) + (10^100 + 3) = 2 × 10^100 + 5
(10^100) + (10^100 + 1) + (10^100 + 2) = 3 × 10^100 + 3
(10^100) + (10^100 + 1) + (10^100 + 3) = 3 × 10^100 + 4
(10^100) + (10^100 + 2) + (10^100 + 3) = 3 × 10^100 + 5
(10^100 + 1) + (10^100 + 2) + (10^100 + 3) = 3 × 10^100 + 6
(10^100) + (10^100 + 1) + (10^100 + 2) + (10^100 + 3) = 4 × 10^100 + 6

以上の10通りあるらしい。
10^100という式がゲシュタルト崩壊して、目と頭がくらくらしてくる。
だがここで、騙されてはいけない。
まともに考えられないようなでかい数が与えられたときは、その数のことは考えなくていいのでは、と疑うことから始める。

冷静になる。Nの制約は最大で200,000ほどなので、例えば全部選んで足し合わせたとしても10^100には到底及ばない。
ということは同じ数字になる可能性があるのは同じ個数を数列を選んだ場合のみ。

つまり、便宜上10^100という数字が与えられているが、ここは０であると考えてもあまり影響はなくて、その時に問題は以下のように言い換えることができる。

「0, 1, 2, 3, 4…NからなるN+1個の等差数列が与えられる。K～N+1個選んだ時、それぞれの場合で作りうる数の個数を全て足す」

すこし考えやすくなってきた。
また、入力例 N = 3 K=2の場合で考えてみる
言い換えた場合の数列は [0, 1, 2, 3]。
まずこの中から２つ選んで作れる数は以下になる。

0 + 1 = 1
0 + 2 = 2
0 + 3 = 3
1 + 3 = 4
2 + 3 = 5

次に３つ選ぶ場合。

0 + 1 + 2 = 3
0 + 1 + 3 = 4
0 + 2 + 3 = 5
1 + 2 + 3 = 6

４つ選ぶ場合は全部選ぶしかない。

0 + 1 + 2 + 3 = 6

一般化して考えてみる。
公差１の等差数列だと、その組み合わせで１番小さい数から１番大きい数まで、必ず作ることができる。
それぞれmin maxとすると(max - min + 1)個の数字が作れる、ということになる。

段々と答えに近づいてきている。
ということはある個数を選ぶ時に、作りうる最大の数から作りうる最小の数が計算できればいい。

ここまでくると、累積和を事前に計算する問題だったんだな、ということが分かる。

入力例は[0, 1, 2, 3]なので、累積和は[0, 1, 3, 6]。
２個選ぶ時の最大は6 - 1 で5。最小は1 。
３個選ぶ時の最大は6 - 0 で6。最小は3。
一般化して、総和をS、累積和の配列Aから i 個選ぶ時の最大、最少は以下のように表せる

max = S - A[n-i]
min = A[i-1]

これをそのまま実装していく。

実装

import { createInterface } from "readline";
import * as fs from "fs";
import { log } from "console";
import { exit } from "process";

function main() {
    const mod = 10 ** 9 + 7
    const [n, k] = nextNums(2)
    // 累積和を構築
    const cumSum: number[] = []
    cumSum.push(0)
    for (let i = 1; i < n + 1; i++) {
        cumSum.push(i + cumSum[i - 1])
    }
    let ans = 0
    const total = cumSum[cumSum.length - 1]
    // kからn個選ぶまで計算
    for (let i = k; i < n + 1; i++) {
        const ma = total - cumSum[n - i]
        const mi = cumSum[i - 1]
        ans += (ma - mi) + 1
        ans %= mod
    }
    // n+1個、つまり全部選ぶ場合を足す
    ans += 1
    log(ans % mod)
}

let inputs = "";
let inputArray: string[];
let currentIndex = 0;

function next() {
    return inputArray[currentIndex++];
}
function nextNum() {
    return +next();
}
function nextBigInt() {
    return BigInt(next());
}
function nexts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = next();
    return arr;
}
function nextNums(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextNum();
    return arr;
}
function nextBigInts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextBigInt();
    return arr;
}


// デバッグ環境がWindowsであれば条件分岐する
if (process.env.OS == "Windows_NT") {
    const stream = createInterface({
        input: process.stdin,
        output: process.stdout,
    });
    stream.on("line", (line) => {
        inputs += line;
        inputs += "\n";
    });
    stream.on("close", () => {
        inputArray = inputs.split(/\s/);
        main();
    });
} else {
    inputs = fs.readFileSync("/dev/stdin", "utf8");
    inputArray = inputs.split(/\s/);
    main();
}

この記事が気に入ったらサポートをしてみませんか？