AtCoder精選良問「D - Handstand」Diff 1138

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

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

問題リンク

D - Handstand

問題概要

N人の人が一列に並んでいる。それぞれ0か１で表されて、0の時は直立していて、１の時は逆立ちをしている。
区間を指定して、1と0を反転させる操作をK回まで行える。
適切に操作を行ったときに、逆立ちしている人（つまり１）を最大で何人並べられるか出力する。

考え方

まずは入力例２を確認してみる。

14 2
11101010110011

2回までの操作で１を何個最大で連続させられるか、というのがこの問題で考えるべきこと。

まず、考察の一つとして、「0の区間だけひっくり返すのが最適」というのがある。
これは考えると当然のことで、例えば1001みたいな場合、間のゼロを２つひっくり返すと、1111になって４つ繋げられる。1を含む形でひっくり返すと、必ず0の区間をひっくり返すよりもパフォーマンスが悪くなる。

つまり、「０の区間をK個までひっくり返して、１を最大で何個続けられるか？」という問題に言い換えることができる。

こういう２種類の文字のみが続く問題は、ランレングスエンコードをして考えるのが、鉄則。
ランレングスエンコードとはaabbという文字があった時に、[[a,2],[b,2]]みたいに圧縮する手法。

入力例を圧縮すると、以下のようになる。

// 11101010110011
[[1,3], [0,1], [1,1], [0,1], [1,1], [0,1], [1,2], [0,2], [1,2]]

操作をすると、０の区間は最大でK個まで１にできる。
つまり、０の区間を最大でK個まで保持して、１をどれくらい続けられるか、を考えればいい。

これには尺取法のアルゴリズムを用いる。

尺取法の解説はいろいろなところにあるので、詳しい説明は省くが、「右端をできるだけ伸ばす。条件を満たさなくなったら、左端を縮める」みたいな感じ。この場合の条件は0の区間がK個以内、ということ。

入力例でいうと、右端を伸ばしていくと、最初の５つまでで０区間がK個、つまり２つになる。

[1,3], [0,1], [1,1], [0,1], [1,1]

この時の長さは、3 + 1 + 1 + 1 + 1 = 7
そこに次の[0,1]という配列がやってくる。
長さを足すと 7 + 1 で8になるが、このとき0を3つ含んでいるので、この長さは答えにすることはできない。

そこで、０がK個に収まるように左端を取り除いていく。
0の区間数を２個以内にするためには、
[1,3], [0,1]の２つの区間を捨てる必要がある。
そのため、長さとしては 8 - 4 で4になる。

このように最大でK個の0の区間を保持して、右端を伸ばしながら左端を捨てていく。過程で計算された長さの最大長を答えればACとなる。

実装

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

/**
 * 
 * @param s 文字列
 * @returns aabb => [['a',2],['b',2]]
 */
export const runLengthEncode = (s: string): Array<[string, number]> => {
    if (s.length == 0) return []
    const encoded: Array<[string, number]> = [];
    let currentChar = s[0];
    let count = 1;

    for (let i = 1; i < s.length; i++) {
        if (s[i] === currentChar) {
            count++;
        } else {
            encoded.push([currentChar, count]);
            currentChar = s[i];
            count = 1;
        }
    }
    encoded.push([currentChar, count]);

    return encoded;
}

function main() {
    const [n, k] = nextNums(2)
    const s = next()
    // ランレングスエンコードする
    const rle = runLengthEncode(s)
    // 答え
    let ans = 0
    // ゼロ区間の個数
    let zeroCount = 0
    // 左端
    let left = 0
    // 1を続けられる長さ
    let length = 0
    for (let right = 0; right < rle.length; right++) {
        const [char, cnt] = rle[right]
        // 0区間だったら足す
        if (char == '0') zeroCount++
        // 長さを足す
        length += cnt
        // ゼロ区間がk個以内になるまで左端を捨てる
        while (zeroCount > k) {
            // 捨てる区間
            const [removeChar, removeCnt] = rle[left]
            // 長さを縮める
            length -= removeCnt
            // 0を捨てたらカウントダウン
            if (removeChar == '0') zeroCount--
            // 左端をインクリメント
            left++
        }
        ans = Math.max(ans, length)
    }
    log(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();
}