ARC074-E RGB Sequence - Python奮闘記

問題

\(N\)要素の列があって、それぞれRGBで塗り分ける。 \(M\)個の束縛条件があって、各条件は区間中の色の種類数を指定する。

気持ち

O(N³) まで許す制約と、条件で束縛される数え上げという見た目から、DP臭がムンムンする。では何を状態と見るかということが問題になる。

条件\(l,r,x\)を読んだときの判定はループが\( r \)に到達したときに考慮するのが良さそう。 3色しかないので、各色について1つずつ状態を持たせられそうである。もっと言えば、各色について最後に使ったindexを状態として持っておけば、何色使ったかを判定することはできる。この条件をちゃんと満たすものだけ遷移させるようにすればOK。

ACコード

Submission #15580318 - AtCoder Regular Contest 074

import sys
input = sys.stdin.readline

mod = 10**9+7

# (i, j, k+1)に置くのが、条件(l,x)に適しているか
def check(i, j, l, x):
    return (x == 3 and l <= i) or (x == 2 and i < l <= j) or (x == 1 and j < l)


N, M = map(int, input().split())
LRX = [list(map(int, input().split())) for _ in range(M)]

LX = [[] for _ in range(N+2)]
for l, r, x in LRX:
    LX[r].append((l, x))

dp = [[[0]*(j+1) for j in range(i+1)] for i in range(N+2)]
# dp[i][j][k] (k <= j <= i) 
# 三色それぞれ、最後に使ったindex(1-indexed)がi,j,k

dp[0][0][0] = 1

for i in range(N+1):
    for j in range(i, N+1):
        for k in range(j, N+1):
            
            # 一番前の色を置く
            ok = True
            for l, x in LX[k+1]:
                if not check(j, k, l, x):
                    ok = False
                    break
            if ok:
                dp[k+1][k][j] = (dp[k+1][k][j] + dp[k][j][i]) % mod
            
            # 最後
            ok = True
            for l, x in LX[k+1]:
                if not check(i, j, l, x):
                    ok = False
                    break
            if ok:
                dp[k+1][j][i] = (dp[k+1][j][i] + dp[k][j][i]) % mod
            
            # 二番目
            ok = True
            for l, x in LX[k+1]:
                if not check(i, k, l, x):
                    ok = False
                    break
            if ok:
                dp[k+1][k][i] = (dp[k+1][k][i] + dp[k][j][i]) % mod
        

ans = 0
for a in range(N+1):
    for b in range(a,N+1):
        ans = (ans + dp[N][b][a]) % mod

print(ans)

感想

DPではありそうとは思ったが、「最後に使ったind」を状態にするのは思いつかなかった。そしてそれを理解した上でもバグらせまくるという、、