【AtCoder】ABC099D Good Grid

問題

https://atcoder.jp/contests/abc099/tasks/abc099_d
簡単に言うと、(i+j)%3の値ごとにグリッドを異なる色で染めるときに最小コストを求める問題

解法1

(i+j)%3は0,1,2のいずれかなので、グリッドの各マスは3種類に分けられる。

3種類のグリッドをそれぞれ1からC(<=30)で塗るときにかかるコストを前計算しておけばあとはグリッドの色の組合せ(30*29*28)を試せばよいのでO(CN^2+C^3)でできる。
(ところでこの計算量の書き方は正しいのかな…)

N,C = map(int,input().split())
D = [list(map(int,input().split())) for _ in range(C)]
c = [list(map(int,input().split())) for _ in range(N)]
 
#cst[s][t]はs色で(i+j)%3がtとなるグループのグリッドたちを塗った時のコストを表す
cst = [[0,0,0] for i in range(C)]

#コストの前計算
for clr in range(1,C+1):
    for i in range(N):
        for j in range(N):
            cst[clr-1][(i+j+2)%3] += D[c[i][j]-1][clr-1]
 
#色の組合せの全探索(同じ色が使われるときはパスする)
ans = 10**10
for i in range(C):
    for j in range(C):
        for k in range(C):
            if i == j or i == k or j == k:
                continue
            ans = min(ans,cst[i][0]+cst[j][1]+cst[k][2])
 
print(ans)

ところがこれだとPyPy3ではACするがPython3だとTLEしてしまう。
解説や他の人のコードを見たところ、もっと早くできるらしい。

解法2

上記の解法の無駄な部分は、全てのマスを1からCの色すべてで塗ることを試しているところで、ここは各グループのグリッドが、それぞれ何色がいくつあるかを前計算しておくことでさらに早くできる。
これだとO(N^2+C^4)でPython3でもTLEしない。

N,C = map(int,input().split())
D = [list(map(int,input().split())) for _ in range(C)]
c = [list(map(int,input().split())) for _ in range(N)]
 
#cst[s][t]は(i+j)%3がtとなるグループに最初s色だったグリッドがいくつあるかを表す
cst = [[0,0,0] for i in range(C)]
 
#個数の前計算
for i in range(N):
    for j in range(N):
        cst[c[i][j]-1][(i+j)%3] += 1

#色の組合せを全探索するのだが、その時にグリッドの種類がlのものすべてをm色からx[l]色に塗り替えるときのコストはD[m][x[l]]*cst[m][l]
ans = 10**10
for i in range(C):
    for j in range(C):
        for k in range(C):
            if i == j or i == k or j == k:
                continue
            x = [i,j,k]
            a = 0
            for l in range(3):
                for m in range(C):
                    a += D[m][x[l]]*cst[m][l]
            ans = min(ans,a)
 
print(ans)

グリッドの位置はどうでもよくて、どのグループに所属しているかが大切なのでまとめられると考えれば個数を前計算するというのは思いつけそう。
個数を前計算してまとめておくのは結構汎用性が高そうだと思った。