[AtCoder]ABC B065 – Trained?[C++]

引き続き、AtCoderのB問題をC++で解いていきます。

ABC B065 – Trained?

問題

筋力をつけたい高橋君は、AtCoder社のトレーニング設備で、トレーニングをすることにしました。
AtCoder社のトレーニング設備にはN個のボタンがついており、ちょうど1個のボタンが光っています。 ボタンには、1からNまでの番号がついています。 ボタンiが光っているときにそのボタンを押すと、ボタンiの明かりが消え、その後ボタンaiが光ります。i=aiであることもあります。 光っていないボタンを押しても、何も起こりません。
最初、ボタン1が光っています。高橋君は、ボタン2が光っている状態で、トレーニングをやめたいと思っています。
そのようなことは可能かどうか判定し、もし可能なら最低で何回ボタンを押す必要があるかを求めてください。

自分の回答

#include <bits/stdc++.h>
using namespace std;
 
int main () {
  int N;
  cin >> N;
  // 明かりがついているか
  // 1なら点灯している、0なら消えている
  int light[N];
  for(int i = 0; i < N; i++){
    light[i] = 0;
  }
  light[0] = 1;
 
  // ボタンの情報
  int btn[N];
  for(int i = 0; i < N; i++){
    int a;
    cin >> a;
    btn[i] = a;
  } 
  //次に押すボタン、最初はボタン1を押すので、初期値は1
  int next = 1;
  
  for(int i = 0; i < N; i++){
    // ボタンが光っている場合
    if(light[next - 1] == 1){
      next = btn[next - 1];
      light[next-1] = 1;
      if(next == 2){
        cout << i + 1<< endl;
        exit(0);
      }
    // ボタンが光っていない場合
    }else{
      next = btn[next - 1];
    }
  }
  cout << "-1" << endl;
}

少し難しかったです。

まずは、条件として、光っているボタンを押すと、次のボタンが光り、光っていないボタンを押しても何も起きないということなので、ボタンの点灯状態を記録する為の、配列lightを用意します。もし、値が1の時は光っている、0の時は光っていないこととします。最初はボタン1以外は光っていないので、初期値0を代入しておきます。ボタン1だけは、最初の状態から光っているので、light[0] = 1としています。
続いて、次にボタンの情報を代入する配列btnを用意します。次に押すボタンを変数nextに代入するようにします。最初はボタン1を押すので、初期値はnext=1となります。
あとは、ボタンが光っているかどうかを判断していけばいいです。ボタンが光っている場合は、変数nextに次に押すボタンの値を渡して、その値のボタンを光らせます。そして、次のボタンが2だった場合は、そこでボタンを押した回数を出力して終わりになります。
次のボタンの明かりがついていなかった場合、何も起こらないので、次に押すボタンの値を代入するだけです。