Python3で解く AtCoder Beginner Contest 175 D - Moving Piece

当たり前だが、制約が大きいので、一つ一つ移動させてスコアの和の最大値を求めようとするとTLEになる。そこを工夫する必要があり、 $2^i$ の和で考えて対応した。最初に選んだマスがある意味運命を左右することになるので（止まるという選択肢はあるものの）そこをうまく実装する必要がある。インデックス操作で頭が大混乱するのでそこらへんが注意ポイント。

概要
解くときに考えた内容
コード
書籍
- 最近ポチった書籍
- アルゴリズムの参考書籍

概要

問題
1. $1$ から $n$ 番までの番号のついた $n$ マス目上でコマを動かすゲーム考える。
1. $i$ 番目のマス目には整数 $c_i$ が書かれている。のちに説明するが、これはこのマス目のスコアである。
1. また、 $i$ 番目のマス目に対応している、 $p_i$ が与えられる。
1. 1回の移動では $p_i$ は現在 $i$ 番目( $1 \leqq i \leqq n$ )のマスにいるなら、次は $p_i$ 番目のマスへコマを移動させる。
1. 移動後、 $c_{p_i}$ がスコアに加算される。
1. $1$ から $k$ 回動ける。
1. ゲーム終了時のスコアとしてあり得る値の最大値は？（ゲーム開始時のスコアは 0）

解くときに考えた内容

まずと制約が大きいので、１つ１つ移動させるのではなく、2の累乗の形に落とし込む。
- $10^9$ を $2^i$ の和で表せば計算量を大きく削減できる。
- 例えば、10回動くというのは $2^3 + 2^1$ というようなイメージ
- 例えば、100回動くというのは $2^6 + 2^5 + 2^2$ というようなイメージ
- 2進数のbitで考えるようなイメージ
- 制約が $10^9$ なので $2^{31}$ まであれば十分
今回のキモは1回でも動けば動かなくてもよくて、その上でのスコアの最大値を求めなくてはいけないということ。
大事なのは最初に選択したマス $j$ で、それ以降は指示通りに動くしかないので、最初に決めたマスといつまで続けるかだけを判断できるようにすればいいということ
いくつかの $n \times 31$ の多重リストを作る。
- $p$ を扱いやすいようにindex化する（左から何番目 - 1）
- posi : $posi_{ij}$ は、最初に選択したマスを $j$ で表す。0なら左から1つ目のマス、1なら左から２つ目のマス。 $i-1$ の時に指示通りに $2^{i-1}$ 回移動したindex。初期値は $p_j$
- point : $point_{ij}$ は、最初に選択したマスを $j$ で表す。0なら左から1つ目のマス、1なら左から２つ目のマス。 $i-1$ の時に指示通りに $2^{i-1}$ 回移動した時に獲得する累計スコア。初期値は $C_{p_j}$
- point_max : $point\_max_{ij}$ は、最初に選択したマスを $j$ で表す。0なら左から1つ目のマス、1なら左から２つ目のマス。 $i-1$ の時に、指示通りに $2^{i-1}$ 回移動した時に獲得する累計スコアの「最大値」pointとの違いは、低くなるような値を足さないということ（このリストを作れるかどうかがキモ）
場合分けとして、
- スコアが全部負であれば、 $c_{p_i}$ の最大値を選ぶ。
- スコアに正の数が含まれているなら、posiとpointとpoint_maxを作っていく。
多重リストの一つ手前の値を参照しながら作る。
- 初期値はindexが最初に選んだ場所で1回動いた時の状態にする。
- そうすると、各リストのindexを $i$ とすると最初から $2^i$ 回動いた時点での値になる。
- 理由は単純で、累積和になるので、一つ手前の状態で、一つ手前の動きをするから 1, 2(1+1), 4(2+2), 8(4+4), 16(8+8)動いた状態であるということになる。
- posiはルール上、動くなら今いる場所に書いてあるマスにしか動けないのでそのように書く
- pointは、一つ手前の状態のpoint + 一つ手前の状態のpointで最初に選択したマス $j$ に書かれているposiのpointを足す。（頭が混乱するので注意!!）
  posiとpointのイメージ(入力例1を参考にしている)
- point_maxはさらに混乱するが $point\_max_{ij}$ を計算する時には、 $posi_{i-1}$ から $2^{i-1}$ 回動くときに、負の値を無駄に足すことがないようにする。
- コードを見る方が早いが、point_maxとしては、現状維持（つまり、point_maxの一つ前の値）と一つ前のpointと最初に選択したマス $j$ に書かれているposiの一つ前のpoint_maxを足したもののうち大きい方を選ぶようにする。（超大混乱するので、printしながら確認するとよい）そうでないと、無駄にマイナスされてしまうため。ある状態になったらこれ以上スコアを下げないという処理をmaxをとることで、point_maxで処理する。
ここまでくれば、あとは $k$ が2の何乗かがわかれば良い。
- 今回はbit演算で行った。
- 単純に割り算していく方法もコメントアウトしておく
最後は
- 2パターン。１つはの形で表されるで表される場合。
  - この時は最初に選んだマスでどれが一番大きい値かをチェック
もう一つは、 $k = 2^a + 2^b$ のように複数の組み合わせになっている場合。この時は、最初に選んだマスと、 $k = 2^a$ が回りきった時のマスに書かれている指定のマスから始めなくてはいけないのでその処理をする。

f:id:black_tank_top:20200819094639j:plain — 最後のpoint_maxの処理のイメージ(入力例1を参考にしている)

お疲れ様!!　時間内にこれをまとめきることは厳しい。。。まだまだ精進が足りない。

コード

n, k = map(int, input().split())
p = [int(x) - 1 for x in input().split()]
c = [int(x) for x in input().split()]
pow_num = 31
point = [[0] * n for i in range(pow_num)]
point_max = [[0] * n for i in range(pow_num)]
posi = [[0] * n for i in range(pow_num)]

# 全部スコアが負だったら一番大きいものを出力する
c_set = set(c)
if all([i < 0 for i in c_set]):
    print(max(c_set))
# スコアに正の数が含まれているなら以下のようにやる
else:
    # posi:indexは最初に選んだindex（左からのマス番目-1）で、そこから2**i回移動した時のindex
    # point:indexは最初に選んだindex（左からのマス番目-1）で、そこから2**i回移動した時に獲得できる累計point
    # 2**30 > 10**9なので
    for i in range(pow_num):
        if i == 0:
            # 初期値設定。
            for j in range(n):
                posi[i][j] = p[j]
                point[i][j] = c[p[j]]
            continue
        for j in range(n):
            posi[i][j] = posi[i-1][posi[i-1][j]]
            point[i][j] = point[i-1][j] + point[i-1][posi[i-1][j]]
    # 状況判断用
    # print(posi)
    # print(point)
    # 仮にある場所から2**i回移動しても負の値を無駄足すことのないようにする
    # 例えば4回移動するけど、4回までの和よりも途中の和の方が大きいならそこで止める
    for i in range(pow_num):
        if i == 0:
            for j in range(n):
                point_max[i][j] = point[i][j]
            continue
        for j in range(n):
            # 止まるか、動くにしても一番累積の高いスコアを足すかで
            point_max[i][j] = max(point_max[i-1][j], \
                point[i-1][j] + point_max[i-1][posi[i-1][j]])
    # 状況判断用
    # print(point_max)
    # kが2の何乗の和であるかっていう組み合わせリストにする
    # k = 10だったら2**3 + 2**1 なので[3, 1]
    # bitで考えるなら
    power_count = list()
    for i in range(pow_num):
        # kを2進数で考えた時の右から1が立っているかどうか
        if k & (1 << i):
            power_count.append(i)
    # # 単純に割り切れるかどうかだけを見るなら
    # # power_countは以下のように求めてもいい
    # _power = 0
    # while k:
    #     if k % 2 == 1:
    #         power_count.append(_power)
    #     k //= 2
    #     _power += 1
        ans = 0
        for i in range(n):
            next_posi = i
            cumsum_point = 0
            # k == 2**iの形で表されるなら、n回まわるだけ
            # そうでないなら、 i=0から始まって、power_count回す時に
            # 次の遷移先はposiから引っ張ってくる
            for posi_idx in power_count:
                if ans < cumsum_point + point_max[posi_idx][next_posi]:
                    ans = cumsum_point + point_max[posi_idx][next_posi]
                cumsum_point += point[posi_idx][next_posi]
                next_posi = posi[posi_idx][next_posi]

        print(ans)

書籍

アルゴリズムの参考書籍

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造

いわゆる、螺旋本。Aizu Online Judgeの問題を題材に解説されている。自分としては、 Aizu Online Judgeは最近はβ版運用中で、構成がアルゴリズムの項目ごとに問題がタグ付けされているようなイメージで、ここら辺弱いから集中的にやってみようっていう感じでできるので、いいと思っている。（AtCoderもそんな感じのリスト作って欲しい）ただ、この書籍はC++で書かれているのでPythonでやる場合はあくまでもその過程の考え方を参考にし自分で実装する必要がある。（AOJも解いているので、記事にしていこうかなぁ）