AtCoder ABC 175 反省と D問題について

atcoder.jp

結果

A, B, C の3完だった。A と C は凡ミスで一度 WA になってしまった。残念。

D 問題の復習

問題文

記事最上部のリンク参照。

考え方

コマの移動の仕方を有向グラフに起こすと分かりやすいと思う。各マスの移動先は順列で与えられるため、グラフの頂点は出ていく辺と入ってくる辺を一本ずつ持つ。ゆえにグラフは1 つ以上のサイクルとなる。
例えばサンプル1の場合は次のような有向グラフとなる。

始点の選びかたと移動の回数を全探索して答えを求める。
まずは各サイクル上で適当な始点を選んだとき、1回ループまでの間に獲得する点数の最大値を計算する。ただし、K が選んだ始点の所属するループの長さより短い場合は、1K の間で最大の値を計算することになる。K がループの長さより大きい場合は一度サイクル上をループする間に獲得する点数の大きさに応じて、場合分けをする必要がある。

  • 一度ループする間に獲得する点数の合計が 0 以下場合

この場合はループすればするほど点数を損するため、先ほど計算した値が最大値となる。

  • 一度ループする間に獲得する点数の合計が 0 以上場合

この場合はループすればするほど点数を得することができる。よって、先ほど求めた値に、可能な限りループを繰り返すことによって得られるポイントを加えたものが最大値となる。

コード
#include<bits/stdc++.h>
using namespace std;
int main(){
  int n,k;cin>>n>>k;
  vector<int>p(n),c(n);
  for(int i=0;i<n;i++){
    cin>>p.at(i);
    p.at(i)--;
  }
  for(int i=0;i<n;i++){
    cin>>c.at(i);
  }
  vector<vector<int>>G(n,vector<int>());
  int64_t ans=-1e18;
  for(int i=0;i<n;i++){
    int x=i;
    int64_t sum=0;
    int64_t tmp=0;
    while(1){
      G.at(i).push_back(c.at(x));
      sum+=c.at(x);
      x=p.at(x);
      if(x==i)break;
    }
    for(int j=0;j<k;j++){
      int roop_size=G.at(i).size();
      if(j>roop_size-1)break;
      tmp+=G.at(i).at(j);
      if(sum>0){
        ans=max(ans,tmp+sum*((k-j-1)/roop_size));
      }
      else{ans=max(ans,tmp);}
    }
  }
  cout<<ans<<endl;
}