C++でn日後の木の高さを求めるアルゴリズムと実装コード

問題概要

n個の要素を持つ配列Aがあるとします。Aの各要素は0または1の値を取ります。ここに一本の木があり、連続するn日間について、A[i]が0の場合はその日に水やりを行わず、1の場合は水やりを行うものとします。木は以下のルールに従って成長します。

• 2日連続で水やりが行われないと、木は枯れてしまいます
• i日目に水やりをすると、木は1cm成長します
• i日目と(i+1)日目に連続して水やりをした場合は、1cmの代わりに5cm成長します
• i日目に水やりをしなかった場合は、木は成長しません

最初、木の高さは1cmです。n日後の木の高さを求めてください。ただし、木が枯れてしまった場合は-1を返します。

例えば、入力がA = [0, 1, 1]の場合、出力は7になります。これは以下のように計算できます。

• 1日目：水やりなしのため成長せず、高さは1cmのまま
• 2日目：水やりをして1cm成長し、高さは2cmに
• 3日目：前日に続き水やりをしたので5cm成長し、2 + 5 = 7cmになる

解法のアプローチ

この問題を解くには、以下の手順に従います。

1. 現在の高さを表す変数rを1、前日の水やり状態を表す変数yを0で初期化します
2. 配列Aの各要素xを先頭から順に処理します
3. rがすでに-1（木が枯れている）の場合は、以降の処理をスキップします
4. xとyがどちらも非ゼロ（連続水やり）の場合は、rに5を加算します
5. xのみが非ゼロの場合は、rに1を加算します
6. xとyがどちらもゼロで、かつi > 0（初日以降）の場合は、2日連続無水のためrを-1に設定します
7. 最後にyへxを代入し、翌日の処理に備えます

r := 1
y := 0
n := size of A
for initialize i := 0, when i < n, update (increase i by 1), do:
    x := A[i]
    if r is same as -1, then:
        Ignore following part, skip to the next iteration
    if x is non-zero and y is non-zero, then:
        r := r + 5
    otherwise when x is non-zero, then:
        (increase r by 1)
    otherwise when not x is non-zero and not y is non-zero and i > 0, then:
        r := -1
    y := x
return r

実装例

それでは、理解を深めるために実際のC++での実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
int solve(vector<int> A){
    int r = 1;
    int y = 0;
    int n = A.size();
    for (int i = 0; i < n; ++i){
        int x = A[i];
        if (r == -1)
            continue;
        if (x && y)
            r += 5;
        else if (x)
            ++r;
        else if (!x && !y && i > 0)
            r = -1;
        y = x;
    }
    return r;
}
int main(){
    vector<int> A = { 0, 1, 1 };
    cout << solve(A) << endl;
}

入力

{ 0, 1, 1 }

出力

7

まとめ

このアルゴリズムは、配列を一度だけ走査するO(n)の計算量で動作します。前日の水やり状態を記憶しておくことで、連続水やりによる5cm成長と、2日連続無水による枯死の判定を効率的に行えるのがポイントです。同様のシミュレーション問題は競技プログラミングでも頻出のため、このパターンは覚えておくと役立ちます。