AtCoder精選良問「B - Count 1's」Diff 883

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

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

問題リンク

B - Count 1's

問題概要

１と０からなる整数の配列が与えられた時に、1つの操作を行うことで得られる「スコア」の取りうるパターン数を求めるという問題。

操作：
配列の連続する部分列を選んで、その中の要素を全て反転させる（0を1に、1を0に変える）。
選ぶ部分列が空であっても構わず、その場合配列の要素は変化しない。

ここで、「スコア」は配列に含まれる1の数を指す。

考え方

結構難しい。ナイーブな解法としては、左端と右端を全探索することだが、制約が300,000まであるので、厳しいか。効率的に考える必要がある。
結論から言うとkadane's algorithm(最大部分配列和)を使って解くことができる。

まずは入力例１をみてみる。

4
0 1 1 0

何も操作しないことも可能なので、上の例でいくと２というスコアは確実に作れるということはすぐに分かる。

ここからどれだけスコアを増減させることができるか、ということを考えていくとどうだろう。

まずスコアを増やせるということは選んだ部分列において、０の数が１より多いということになる。
逆にスコアを減らせるということは１の数が０の数より多い。

ということは、（０の数）－（１の数）の最大数がとりうる最大増加。
逆に（１の数）ー（０の数）の最大数がとりうる最大減少。

紹介したkadane's algorithmは配列の最大部分列和を求めるもの。
関数にすると以下のような形になる。

/**
 * 
 * @param arr 配列
 * @returns 最大部分配列和
 */
export const maxSubarraySum = (arr: number[]): number => {
    let maxCur = arr[0]
    let maxGlobal = arr[0]
    for (let i = 1; i < arr.length; i++) {
        maxCur = Math.max(arr[i], maxCur + arr[i])
        if (maxCur > maxGlobal) {
            maxGlobal = maxCur
        }
    }
    return maxGlobal
}

このアルゴリズムは正負の数が混じっている配列の最大部分列和を教えてくれる。

[1, -3, 2, 1, -1]

例えばこんな配列だったら、途中の２＋１が最大なので、３を返す。

これを使って考えてみる。

0 1 1 0

入力例の配列はこんな感じだった。
まず、０を選んだ時はスコアが増える、逆に１の時はスコアが減るので、０のときは１、１の時は-1としてみよう。

1 -1 -1 1

こうなる。
これの最大部分列和は１。つまり１だけ数を増やすことができる。

逆に減らせる数はこれを反転させて考える。
０のときはー１、１のときは１。

-1 1 1 -1

これの最大部分列和は２。つまり２だけスコアを減らせる。
増やす場合も減らす場合も、必ず１ずつしか増減しない。

そのため最大値ー最小値の区間がとりうる答えとなる。

この場合だと、元の数は２。
そのため０(2-2)～３(2+1)の４つのスコアを取ることができる。

実装

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

/**
 * 
 * @param arr 配列
 * @returns 最大部分配列和
 */
export const maxSubarraySum = (arr: number[]): number => {
    let maxCur = arr[0]
    let maxGlobal = arr[0]
    for (let i = 1; i < arr.length; i++) {
        maxCur = Math.max(arr[i], maxCur + arr[i])
        if (maxCur > maxGlobal) {
            maxGlobal = maxCur
        }
    }
    return maxGlobal
}

function main() {
    const n = nextNum()
    const arr = nextNums(n)
    // 元の数
    let base: number = 0
    for (const num of arr) {
        base += num
    }
    // 最大で増やせる数の計算用
    const arrIncrease: number[] = []
    // 最大で減らせる数の計算用
    const arrDecrease: number[] = []
    for (const num of arr) {
        if (num == 1) {
            arrIncrease.push(-1)
            arrDecrease.push(1)
        } else {
            arrIncrease.push(1)
            arrDecrease.push(-1)
        }
    }
    // 増やせる最大量
    const increaseMax = Math.max(0, maxSubarraySum(arrIncrease))
    // 減らせる最大量
    const decreaseMax = Math.max(0, maxSubarraySum(arrDecrease))
    // 取りうる範囲を計算
    const ansMax = base + increaseMax
    const ansMin = base - decreaseMax
    // 答えの出力
    log(ansMax - ansMin + 1)
}

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