AtCoder Beginner Contest 175 D,F問題メモ

C, D問題は、漏れなく考察できてるかの確信を見誤りやすい。

D - Moving Piece

問題

$1～N$ の番号が付いたマス上でゲームを行う
- マスには $C_1,C_2,...,C_N$ の数字が書かれている
- 各マスからの移動先を示す $1～N$ の順列 $P_1,P_2,...,P_N$ が与えられる
ゲームは、まず好きなマスにコマを置き、以下の操作を繰り返す
- 今いる位置を $i$ として、$P_i$ に駒を動かす
- 動かした先に書かれている数字（$C_{P_i}$）をスコアに加算する
操作を $1$ 回以上 $K$ 回以下行ったときのスコアの最大値を求めよ
$2 \le N \le 5000$
$1 \le K \le 10^9$
$P_i \neq i$
$-10^9 \le C_i \le 10^9$

解法

「順列 $P$ があり、$i→P_i$ への移動を繰り返す」という操作は、いくつかの要素からなるサイクルをぐるぐる回ることになる。

i   1   2   3   4   5
P   2   4   5   1   3
C   3   4 -10  -8   8

1→2→4→1→...
3→5→3→...

$O(N^2)$ 解法

各マスからスタートしてシミュレートする。

$K$ がでかいので最後までやるとTLEだが、最大でも $N$ 回たどるとスタートに戻ってくるので、サイクルの長さ $l$ がわかったら周回数 $d$ と余り $m$ を計算する。

             1 → 2 → 4 → 1 → ...
加算スコア      4   -8    3    4 ...
累計スコア      4   -4   -1    3

K=7 l=3  →  周回数=2  余り=1

1周の累積スコアが正
- $d$ 周ギリギリまで回って「周回数×1周のスコア」を稼ぎ、残り $m$ 回を改めてシミュレートする
- これは嘘で、後述のようなケースがある
1周の累積スコアが負や0
- 1周するまでの最大値で終了した方がよい

ただし、1周の累計スコアが正でも、最後の1周に関しては途中で終了するのが最良のケースがある。

K=12       ┌------サイクル--------┐
加算スコア  -1  100    1  100 -199  |  -1  100    1  100  -199  |  -1  100    1  100  -199  ...
累計スコア  -1   99  100  200    1  |   0  100  101  201     2  |   1  101  102  202     3  ...

この場合、回数的には2周でき、1周のスコアも正なのだが、2周後の余り操作2回の中で到達できる累計スコアの最大値は $101$ である。

しかし、前の周回の途中では $201$ を達成できていて、真の最大値はそれである。

従って、1回以上の周回が可能な場合においては、周回数を1回減らし、最後のループを含めて辿るようにするとよい。

加算スコア  -1  100    1  100 -199   |  -1  100    1  100  -199  |  -1  100    1  100  -199  ...
累計スコア  -1   99  100  200    1   |   0  100  101  201     2  |   1  101  102  202     3  ...
                                                  ここからでなく |-------->
                ここからシミュレート |------------------------------------>

$O(N \log{N})$ 解法

一度求めたサイクルの累計スコアの配列は、同じサイクルに属する各マスをスタートした場合で使い回せる。

$i→P_i$ を辿っていってサイクルを発見したら、3周ほどの累計スコアを計算する。先頭に便宜的に0を追加する。

acc   0  -1   99  100  200   1   0  100  101  201   2   1  101  102  202   3

周回後の残り操作回数（1回以上のループが可能な場合は最後の1ループを含めて）を $m$ とすると、$i$ からスタートしたときの最大スコアは、以下で求められる。

acc   0  -1   99  100  200   1   0  100  101  201   2   1  101  102  202   3
         ~~~~~~~ m個の中の最大値 ~~~~~~

max(acc[1:1+m]) - acc[0]  +  1周のスコアが正の場合は(周回数x1周のスコア)

配列 $A$ の区間和を累積和 $Acc$ を使って高速に求めるテクニックとして、$A_l+A_{l+1}+...+A_{r}=Acc[r]-Acc[l-1]$ とするのはよく出てくる。

今回の場合、開始地点 $l$ を決めると、$r$ は $l～l+m-1$ の中から好きに選べる。その中で $Acc[r]-Acc[l-1]$ を最大化するには、$Acc[l]～Acc[l+m-1]$ の区間最大値を求めればよい。

次、$i$ から1回移動した $P_i$ からスタートしたときの最大スコアは、1つずらして以下で求められる。

acc   0  -1   99  100  200   1   0  100  101  201   2   1  101  102  202   3
              ~~~~~~~ m個の中の最大値 ~~~~~~

max(acc[2:2+m]) - acc[1]  +  1周のスコアが正の場合は(周回数x1周のスコア)

平衡二分探索木またはそれに類似したデータ構造を使って区間最大値を管理しておけば、1回の更新に $O(\log{N})$ で済むので、全体で $O(N \log{N})$ で求められる。

from collections import Counter
from heapq import heapify, heappop, heappush
from itertools import accumulate


def solve(n, k, ppp, ccc):
    NINF = -(10 ** 18)
    ans = NINF

    checked = [False] * n
    for s in range(n):
        if checked[s] == True:
            continue
        checked[s] = True
        scores = [ccc[ppp[s]]]
        v = ppp[s]
        while v != s:
            scores.append(ccc[ppp[v]])
            checked[v] = True
            v = ppp[v]

        l = len(scores)
        d, m = divmod(k, l)
        loop = sum(scores)

        if d > 0:
            d -= 1
            m += l

        scores += scores * 2
        scores.insert(0, 0)
        acc = list(accumulate(scores))

        base = max(0, loop * d)

        range_max = [-a for a in acc[1:m + 1]]
        available_max = Counter(range_max)
        heapify(range_max)
        for i in range(l):
            while available_max[range_max[0]] == 0:
                heappop(range_max)
            ans = max(ans, base - range_max[0] - acc[i])
            old = -acc[i + 1]
            new = -acc[i + m + 1]
            available_max[old] -= 1
            available_max[new] += 1
            heappush(range_max, new)

    return ans


n, k = map(int, input().split())
ppp = list(map(int, input().split()))
ppp = [p - 1 for p in ppp]
ccc = list(map(int, input().split()))

print(solve(n, k, ppp, ccc))

E - Picking Goods

問題

解法

何個拾ったかを状態に加える素直なDP。PythonではNumbaでようやく間に合った。

$R \times C \times 4$ のDP配列を作るのだが、$R$ については1行前の情報以外は不要で下から計算できるので、配列を $C \times 4$ として逐次更新していけば多少の高速化になり、PyPyでもACできる、かも。

言語アップデートでNumbaが使える前は、「ループ内で dp[i] に繰り返しアクセスするときはループ外で dpi = dp[i] と変数に入れておくと少し高速化される」とかやっていたが、やっぱり可読性は低まるし、そういうことをせずNumbaで書けば間に合うようになったのは有難い。

import os
import sys

import numpy as np


def solve(inp):
    r, c, k = inp[:3]
    items = np.zeros((r + 1, c + 1), dtype=np.int64)
    rrr = inp[3::3]
    ccc = inp[4::3]
    vvv = inp[5::3]
    for r_, c_, v in zip(rrr, ccc, vvv):
        items[r_][c_] = v

    dp = np.zeros((r + 1, c + 1, 4), dtype=np.int64)

    for i in range(1, r + 1):
        for j in range(1, c + 1):
            up = dp[i - 1][j][3]
            for w in range(4):
                dp[i][j][w] = max(dp[i][j - 1][w], up)
            v = items[i][j]
            if v == 0:
                continue
            for w in range(2, -1, -1):
                dp[i][j][w + 1] = max(dp[i][j][w + 1], dp[i][j][w] + v)

    return dp[-1][-1][3]


if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', '(i8[:],)')(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit('(i8[:],)', cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

F - Making Palindrome

問題

$N$ 個の、英小文字からなる文字列 $S_1,S_2,...,S_N$ と、そのコスト $C_1,C_2,...,C_N$ が与えられる
これらの文字列から1つ以上選び、連結させて回文を作りたい
- $S_i$ を利用するには $C_i$ のコストがかかる
- 同じ文字列は何度でも使ってよいが、コストはそのたびにかかる
何でもいいので回文を1つ作るのにかかる最小コストを求めよ
不可能な場合は-1を出力せよ
$1 \le N \le 50$
$1 \le |S_i| \le 20$
$1 \le C_i \le 10^9$

解法

最短経路探索に帰着。ただ、どのようなデータを用意しておけばいいかの整理がなかなか厄介。

まず、中心となり得る文字 or 2文字を洗い出す。そして、そこを中心とした場合に前・後ろにどんな文字列が余るかを経路探索のノードとし、ダイクストラなどで遷移していく。

具体的には、経路探索におけるノードは、以下の情報を持つ。

前に余ってるか後ろに余ってるか
余っている文字列

始点を求める

たとえば、以下の文字列があると、始点ノードは次のようになる。

abccbadab

中心     余り方向  余り文字列
aの手前  後        abccbadab
a        後        bccbadab
cc       後        dab
d        前        abcc
b        前        abccbada
bの直後  前        abccbadab

これは、地道に各文字を中心として確認していっても、文字列1つあたり $O(|S_i|^2)$ で全て洗い出せる。

ただし、最初から単独で回文となっている場合は、それ1つを選んだときのコスト（複数ある場合はその最小値）を暫定コストとする。

遷移

あるノードに訪れたとき、次のノードへは以下のように遷移する。

$N$ 個全ての文字列について、余りを解消する方に繋げられるかを確認し、可能なら次に余る文字列を求める

たとえば、「余り方向: 後、余り文字列: dab」から次のノードへの遷移を求める。

現在の暫定文字列はこんな感じになっている
  ◎☆◆...◆☆◎dab

余り方向が'後'の場合、それを解消すべく繋ぐのは'前'で、以下の文字列を繋げられる。

ズバリ bad という文字列
末尾が …bad の文字列
d, ad など、全体が bad の接尾辞である文字列

繋げる文字列を $S_i$ として、

1.の場合、回文を作れる。これ以上続けて繋げる意味は無いので、暫定スコアを更新する。

2.の場合、次ノードは「余り方向: 前、余り文字列: $S_i$ の末尾から bad を取り除いた文字列」となる。

3.の場合、次ノードは「余り方向: 後、余り文字列: dab の先頭から $S_i$ の逆を取り除いた文字列」となる。

余り方向が前の場合も同様に整理する。

高速化

全文字列の長さの総和を $T$ とすると、ノード数は、方向が2通り、文字列がどこより前/後かで $|S_i|$ 通り、全文字列合わせて $2T$ 個。
辺数は各ノードからそれぞれ $N$ 個の文字列を確認するので、$2TN$ 本。
Dijkstraの計算量は $O((E+V)\log{V})$ なので、最大値を当てはめると約 $5 \times 10^6$ となる。

実際には、各文字列がどれも長く、かつほぼ全ての接頭・接尾文字列がことごとく余り文字列となるケースはまず無い（と思う……）ので、この数分の一となると思われる。

これだけだと行けそうだが、上記の解法を文字列のまま行おうとすると、遷移時に次の文字列を繋げられるかのチェック・繋げられる場合のノードの作成に、文字列の切り出し・反転・照合という $O(|S_i|)$ 程度の計算量が毎回かかる。

これを含めると無理な可能性も高いので、高速化したい。

一方法として、「各文字列の $j$ 文字目から $k$ 文字目まで」を、反転した場合も含め、全てローリングハッシュ等で整数化しておく。

すると、次に繋げる文字列との照合は「$S_i$ の $p$ 文字目から $q$ 文字目のハッシュ」と「$S_j$ の $r$ 文字目から $s$ 文字目のハッシュ」が等しいか？という配列参照と整数値比較で計算できるので、$O(1)$ で行えるようになる。

事前計算も、時間・空間共に $O(N|S_{max}|^2)$ で、十分小さい。

ただし、複数の文字列長が混在するローリングハッシュで照合を行う場合、$a=0,b=1,...$ とするとたとえば b, ab, aab, … などが全て'1'で同一となってしまう。これを避けるには、「何もない文字」を0として、$a=1,b=2,...$ と対応づけることで、文字列長の違いも区別できる。

またハッシュには衝突がつきものだが、今回は文字列長が短く $27^{20} \simeq 10^{28}$ 程度なので、Pythonならmodで割らずそのまま持てば衝突の心配が無く安心できる。

from heapq import heapify, heappop, heappush


def rolling(s):
    """
    文字列 [i, j] のハッシュ
    j < i の場合は逆順
    """
    n = len(s)
    ret = [[0] * n for _ in range(n)]

    for i in range(n):
        rh = 0
        for j in range(i, n):
            b = ord(s[j]) - 96
            rh = rh * 27 + b
            ret[i][j] = rh

    for j in range(n):
        rh = 0
        for i in range(j, -1, -1):
            b = ord(s[i]) - 96
            rh = rh * 27 + b
            ret[j][i] = rh

    return ret


def solve(n, sss, ccc, rhs):
    INF = 10 ** 18
    ans = INF
    q = []
    for i in range(n):
        s = sss[i]
        c = ccc[i]
        l = len(s)

        q.append((c, 0, i, 0))
        q.append((c, 1, i, l - 1))

        for j in range(l):
            k = 1
            while j - k >= 0 and j + k < l:
                if s[j - k] != s[j + k]:
                    break
                k += 1
            else:
                if j - k == -1 and j + k == l:
                    ans = min(ans, c)
                elif j - k == -1:
                    q.append((c, 0, i, j + k))
                else:
                    q.append((c, 1, i, j - k))

        for j in range(l - 1):
            k = 0
            while j - k >= 0 and j + k + 1 < l:
                if s[j - k] != s[j + k + 1]:
                    break
                k += 1
            else:
                if j - k == -1 and j + k + 1 == l:
                    ans = min(ans, c)
                elif j - k == -1:
                    q.append((c, 0, i, j + k + 1))
                else:
                    q.append((c, 1, i, j - k))

    heapify(q)

    checked = []
    for s in sss:
        checked.append([[False, False] for _ in range(len(s))])

    while q:
        cost, fb, i, j = heappop(q)
        if cost >= ans:
            break
        if checked[i][j][fb]:
            continue
        s = sss[i]
        ls = len(s)
        lm = ls - j if fb == 0 else j + 1
        checked[i][j][fb] = True
        rhr = rhs[i][j]

        if fb == 0:
            for ti in range(n):
                t = sss[ti]
                lt = len(t)
                if lm == lt:
                    if rhs[ti][-1][0] == rhr[-1]:
                        ans = min(ans, cost + ccc[ti])
                elif lm > lt:
                    if rhs[ti][-1][0] == rhr[j + lt - 1]:
                        heappush(q, (cost + ccc[ti], 0, i, j + lt))
                else:
                    if rhs[ti][-1][lt - lm] == rhr[-1]:
                        heappush(q, (cost + ccc[ti], 1, ti, lt - lm - 1))
        else:
            for ti in range(n):
                t = sss[ti]
                lt = len(t)
                if lm == lt:
                    if rhs[ti][0][-1] == rhr[0]:
                        ans = min(ans, cost + ccc[ti])
                elif lm > lt:
                    if rhs[ti][0][-1] == rhr[j - lt + 1]:
                        heappush(q, (cost + ccc[ti], 1, i, j - lt))
                else:
                    if rhs[ti][0][lm - 1] == rhr[0]:
                        heappush(q, (cost + ccc[ti], 0, ti, lm))

    if ans == INF:
        return -1
    else:
        return ans


n = int(input())
sss = []
ccc = []
rhs = []
for si in range(n):
    s, c = input().split()
    sss.append(s)
    ccc.append(int(c))
    rhs.append(rolling(s))

print(solve(n, sss, ccc, rhs))

目次