AtCoder精選良問「D - Rain Flows into Dams」Diff 938

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

AtCoder 精選良問集 Diff 茶後半～水前半｜Qlitre｜note

問題リンク

D - Rain Flows into Dams

問題概要

円形に連なる奇数個の山とそれらの間にあるダムが登場する。各山に降る雨がダムに溜まるルールが与えられており、各ダムに溜まった水の量が与えられたとき、各山に降った雨の量を求める問題。

• 円形にN個の山があり、Nは奇数

• 山の間にN個のダムがあり、ダムiは山iと山i+1の間にある (山N+1は山1と同じ)

• 山iに2xリットルの雨が降ると、ダムi-1とダムiにxリットルずつ水が溜まる

• 各山に非負の偶数リットルの雨が降り、各ダムにAiリットルの水が溜まった

• 各山に降った雨の量を求める

考え方

中学生くらいで習うような方程式を考えていくと解ける、というような問題。

競技プログラミングというと、ごりっとしたアルゴリズム、なんとか探索とかなんとか法とかを使って解く問題も多いけど、素朴な実装で解ける問題もある。それが面白いと思った。

まず、以下の入力例を考えてみる。

3
2 2 4

山が３つあって、その間にダムが３つある。それぞれに溜まった水の量が与えられる。

こういう問題は数列だけみても何かを発想するのが難しい。絵に書いてみると発想が浮かぶことがある。仮に山をA,B,Cとしてみた。

かなり下手で申し訳ないが、とにかく絵を描いてみることが大事だ。
そうすると段々と状況が見えてくることもある。

この問題の文脈ではこういうことが起きているんだなということが分かる。

• Aの水の1/2とBの水の1/2を足すと2になる。

• Bの水の1/2とCの水の1/2を足すと2になる。

• Cの水の1/2とAの水の1/2を足すと4になる。

ということはダムの水の量を二倍するとこういう式が成り立つ。

• ①　A + B = 4

• ②　B + C = 4

• ③　C + A = 8

あとは式を変形させていってみる。

• ①　B = 4 - A 

• ②　(4-A) + C = 4

• ②　C = A

• ③　A + A = 8

そのためAが４であることが分かる。
今回の問題の制約上、最後の式の左辺はかならず2Aとなる。
ということは右辺だけ考えればいい。
一般化すると二個目のダムから始めて、どんどん前のダムの水を引いていくイメージ。
この問題でいうと、まず②の４から①の４を引いて０。そして③の８から０を引いて、８という具合だ。

そうして、２で割れば、一番最初の山に流れる水が決まる。
そうすると順番に答えを割り出せる。

実装

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

function main() {
    const n = nextNum()
    // 配列を受け取る
    const arr = nextNums(n)
    // ２倍にする
    const dams: number[] = []
    for (const num of arr) {
        dams.push(num * 2)
    }
    // 前のダムの水を引いていき、最終的に残る量を計算
    let tmp = dams[0]
    for (let i = 1; i < n; i++) {
        tmp = dams[i] - tmp
    }
    // 最初の山に流れる水の量 仮にAとする
    let water = Number(tmp / 2)
    const ans: number[] = []
    for (let i = 0; i < n; i++) {
        ans.push(water)
        // A + B = dams[0]なので　dams[0] - A が次の山の水の量
        water = dams[i] - water
    }
    log(ans.join(' '))
}

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();
}