AtCoder精選良問「D - Snuke Prime」Diff 933

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

https://note.com/kuri_tter/m/m622eeb3d9ca2

問題リンク

D - Snuke Prime

問題概要

株式会社すぬけは、すぬけプライムという支払いプランを提供している。
加入中は1日あたりC円を支払うことで全てのサービスを利用できる。
高橋くんは、この会社のサービスのうちN個を利用する予定。
i番目のサービスはai日目の始めからbi日目の終わりまで利用する。
すぬけプライムに加入していない期間中は、i番目のサービスを利用する際に1日あたりci円を支払う必要がある。

高橋くんが支払う必要のある最小の合計金額を出力する。

考え方

いろいろと書いてあるが、このように単純に言い換えられることに着目をしてみる。

i日目のサービス合計の金額がすぬけプライムの値段より大きかったら、すぬけプライムに加入する。そうでなかったら加入しない。

明らかにこのロジックに沿って支払いを続けるのが最小だ。
難しいのはi日目のサービスの合計金額を求める部分。

例えば極端な例でこんな入力があったとする。

2 2 
1 5 1
2 3 100

1日から5日までのコスト1のサービスと、２日目から３日目までのコスト100のサービスがある。すぬけプライムのコストは2なので、２日目から３日目まではすぬけプライムに加入した方が明らかに良い。この２日目から３日目まで、という区間を検出するのが意外とややこしい。始まりと終わりがあるので、単純にソートするだけでは計算がうまくできないのだ。

こういう問題を解決するものとして、いもす法というアルゴリズムがある。
区間の始まりにプラスの数字を、区間の終わるところにマイナスの数字を順番に打っていく。上の入力例だとこんな感じ。

[1, 100,0,-100,0,-1]

そうして累積和を取るとその日のコストが浮かび上がる。

[1, 101, 101, 1, 1, 0]

あとはこれをすぬけプライムと比較して足し上げていけば最小となる。

これがいもす法の基本的な考え方だ。
しかし、残念ながらこの問題ではこのままでは使えない。
なぜなら日数が最大10**9ととても大きいので、まともにやるとTLEになるからだ。

日数の配列を用意しないでいもす法っぽいすることを考える。
Nの数は2*10**5と比較的に小さいので、以下のような[日数,コスト増減]の二次元配列を作ることを考える。

[[1,1],[2,100],[4,-100],[6,-1]]

このような配列を作っておけばfor文で繰り返せば答えが出せる。

実装

各サービスの日付は重なることはあるので、まずは辞書型でコストの増減を記録して、その後に配列に変換する。答えはとても大きい数になるので、BigInt型で管理する。

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


function main() {
    const [n, primeCost] = nextNums(2)
    // 日付ごとに増減を記録
    const dic: Record<number, number> = {}
    for (let i = 0; i < n; i++) {
        const [a, b, c] = nextNums(3)
        if (dic[a]) {
            dic[a] += c
        } else {
            dic[a] = c
        }
        if (dic[b + 1]) {
            dic[b + 1] -= c
        } else {
            dic[b + 1] = -c
        }
    }
    const event: number[][] = []
    for (const key in dic) {
        event.push([Number(key), dic[key]])
    }
    // 日付が小さい順にソート
    event.sort((a, b) => a[0] - b[0])
    let perCost = 0
    // 大きい数になるのでBigInt型に
    let ans = 0n
    for (let i = 0; i < event.length - 1; i++) {
        const [nowDate, cost] = event[i]
        // コストを足す
        perCost += cost
        // 次のイベントまでの日数
        const period = event[i + 1][0] - nowDate
        // すぬけプライムで払う場合
        if (perCost > primeCost) {
            ans += BigInt(primeCost) * BigInt(period)
        } else {
            ans += BigInt(perCost) * BigInt(period)
        }
    }
    log(String(ans))
}

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