AtCoder精選良問「A - Gold and Silver」Diff 650

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

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

問題リンク

A - Gold and Silver

問題概要

すぬけくんは現在、1グラムの金と0グラムの銀を持っている。
これから彼はN日間で金と銀の取引を行う。
毎日、「何もしない」もしくは「交換する」の2つの行動の中から1つを選ぶ

もし「交換する」を選択した場合、次のことが起こる：

1. Xグラムの金を持っていたなら、全て銀に交換し、X×A_iグラムの銀を得る

2. Xグラムの銀を持っていたなら、全て金に交換し、X/A_iグラムの金を得る

すぬけくんは金の量を最大化するように行動する。
それぞれの日に交換する場合は１，何もしない場合は０を振って出力する。

考え方

このコンテストを主催しているのが大和証券だからか、株の取引みたいな感じがして面白いと思った。

まず入力例１から操作を追ってみる。

3
3 5 2

解説から、２日目と３日目に操作をするのが、最適らしい。

1 日目: なにもしない．
2 日目:1 グラムの金を銀と交換し，5 グラムの銀を得る．
3 日目:5 グラムの銀を金と交換し，2.5 グラムの金を得る．

ナイーブな解法としては全ての交換する・しないを探索すれば答えが出せるが、日数は200,000まであるので、絶望的な計算量だ。

まず重要な操作の一面として、すぬけ君は常に金か銀のどちらかしか持っていない。そのうえで金の数量をどうやったら最適化できるか考える。

単純に考えると、以下の戦略で最適化ができる。

• 銀を一番入手できるときに、金を銀に交換する。

• 金を一番入手できるときに、銀を金に交換する。
  

これはどうすればいいのかというと、上げ基調の頂点で銀を入手する、下げ基調の底で金を入手すればいい。

折れ線グラフを書くと分かりやすい。

赤丸をつけたところで取引をしていくのがかならず最適となる。

次に上げ基調と下げ基調の転換点を調べる方法を考える。
これが結構難しい。

1 3 5 3 1 3 5 

例えばこんな配列があるとする。
前後で比較しても良いが、自分はランレングス表現を使うのが簡単だと思った。

上げ基調をU、下げ基調をDとして、上の配列をこんな風な文字列に変換する。

UUUDDUU

次にランレングス圧縮する。

[[U,3],[D,2],[U,2]]

最初にUを追加するようにすれば、必ずUDUD…と並ぶ。
あとはDにぶつかったら手前のUのインデックスとDのインデックスを処理すればいい。

実装

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 = nextNum()
    const arr = nextNums(n)
    const trend: string[] = []
    trend.push('U')
    for (let i = 1; i < n; i++) {
        if (arr[i] >= arr[i - 1]) {
            trend.push('U')
        } else {
            trend.push('D')
        }
    }
    const trendChar = trend.join('')
    const rle = runLengthEncode(trendChar)
    const stack: number[] = []
    const ans: number[] = new Array(n).fill(0)
    let idx = 0
    for (const elm of rle) {
        const [char, cnt] = elm
        // インデックスを進める
        idx += cnt
        if (char == 'U') {
            // 上げ基調の頂点、stackにインデックスを加える
            stack.push(idx)
        } else {
            // 下げ基調の底。
            // 頂点のインデックスを取り出す
            const peak = stack.pop()!
            // 答えの更新
            ans[peak - 1] = 1
            ans[idx - 1] = 1
        }
    }
    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();
}

ランレングス圧縮しない場合

ランレングス圧縮をしない場合はこんな感じだろうか。

function main() {
    const n = nextNum()
    const arr = nextNums(n)
    // 金を持っているかどうか
    let havingGold = true
    const ans: number[] = new Array(n).fill(0)
    for (let i = 0; i < n - 1; i++) {
        const now = arr[i]
        const nxt = arr[i + 1]
        if (nxt >= now) {
            // 銀を金に換える
            if (!havingGold) {
                ans[i] = 1
                havingGold = true
            }
        } else {
            // 金を銀に変える
            if (havingGold) {
                ans[i] = 1
                havingGold = false
            }
        }
    }
    if (!havingGold) {
        ans[n - 1] = 1
    }
    log(ans.join(' '))
}

コーディング量は減るが、金を持っている・持っていない、上げ・下げの二つの状態を考えるのがちょっとややこしい。実際に書くのに時間がかかった。バグも起きやすい気がする。

ランレングス圧縮して、Dの時だけマーキングする方は分量は増えるが、個人的には思考量が少ないと感じる。引き出しの一つとして持っておきたい。