背景

RECRUIT 日本橋ハーフマラソン 2023夏（AtCoder Heuristic Contest 022）に出てきた、正方グリッド上で隣接マスの値との差分の二乗を最小化する問題を解説します。

問題設定

AHC022では整数$L$と$N$が与えられ、$L\times L$の正方グリッド上のマス目$(i,j)$に整数の温度$P_{i,j}$を設定し、配置コスト

H = \sum_{i=0}^{L-1}\sum_{j=0}^{L-1} \left[ ( P_{i,j}- P_{i,(j+1) \% L} )^2+( P_{i,j}- P_{(i+1) \% L, j} )^2 \right]

を小さくしたい場面がありました。すべての$i,j$に対して同じ温度を選べばコストは最小の$0$になりますが、コンテスト中では温度$P_{i,j}$のうち$N$個を固定し、残りのマスの温度値を動かしてコストを下げたいという場面が多くありました。この記事ではこの問題の近似的な解法を2つ紹介します。いずれもまず温度値を実数であると思って最適化し、最後に整数に直すことで近似値を得ます。

解法1: 平均値で均していく

コストの勾配を取り、それを使って解を改善していくことを考えます。とある${y,x}$に注目したとき、コスト関数に$P_{y,x}$が現れるのは

( P_{y,x}- P_{y,x+1} )^2+( P_{y,x}- P_{y+1,x} )^2 + 
( P_{y,x-1}- P_{y,x} )^2+( P_{y-1,x}- P_{y,x} )^2

の4項ですから、コスト関数の$P_{y,x}$微分は

\begin{align}
\frac{\partial H}{\partial P_{y,x}}
&= 2( P_{y,x}- P_{y,x+1} )+2( P_{y,x}- P_{y+1,x} ) + 
2( P_{y,x}- P_{y,x-1} )+2( P_{y,x}- P_{y-1,x} ) \\
&= 8\left( P_{y,x} - \frac{P_{y,x+1} + P_{y+1,x} + P_{y,x-1}+ P_{y-1,x} }{4} \right)
\end{align}

です。(領域の境界付近ではないことを仮定していますが、今回の設定($L-1$の次が$0$になっている周期境界条件)では境界付近でも同じことが成り立ちます。)ここから最適な配置では$P_{y,x} $の値が周囲4マスの平均値になっていることがわかります。このことから以下を反復することでコストを下げて行くことが可能です。

固定箇所以外の温度を適当に初期化
以下を繰り返す
- $y,x$をランダムや順番に選ぶ。固定マスは選ばないようにするかスキップする。
- $P_{y,x} $の値を周囲4マスの平均値に置き換える。

ちなみに、このアルゴリズムは学習率を$\frac{1}{8}$にした勾配降下法とみることもできます。また問題の設定上温度は整数でなければいけないので、最後に実数値となってしまった温度を一番近い整数値に丸めるとよいでしょう。いずれにしても整数値問題に対して連続値だと思ってアルゴリズムを考えているので、その意味で近似的な解法となっています。

ちなみにこの処理は画像処理の文脈でいう平均値フィルタと言われるもので、scipy.ndimageを使って簡単に実装できます。固定マスを無視した場合はこんな感じで一括処理ができます。固定する箇所の温度値を事前に保持しておけば、numpyの高速性を生かした処理が可能です。

import numpy as np
from scipy import ndimage
kernel = np.array([[0, 1, 0],
                   [1, 0, 1],
                   [0, 1, 0]]) / 4.0
temperature = np.array([[9,1,2],[-2,4,5],[6,7,8]], )

ndimage.convolve(temperature, kernel, mode="wrap")

解法1の計算量

すべてのマスを一回更新するのに$L^2$回の演算が必要で、あるマスの温度の影響が端まで伝搬するには$L$回程度の更新を繰り返す必要があることから計算量のオーダーはおおよそ$L^3$です。今回のコンテストでは$L\sim 100$くらいですので十分実行可能でしょう。更新が小さくなってきたら打ち切ってもよいでしょう。

解法2: 行列演算で解く

温度に関するコスト関数は、すべての項が$P$たちの二次式になっている、二次形式と呼ばれる形をしています。こうした式は$P$たちをうまく並べた縦ベクトル$p$と定数の行列$Q$を使って$p^T Q p$と変形することができます。（さらに今回の問題の場合は行列$Q$を正定値対称に選ぶことができ、高速かつ安定に解けます。）今回のコスト関数の場合は$L^2 $次元のベクトル$p$を、$P_{y,x}$を一列に並べることで定義します。すなわち以下のように成分を決めます。

p[i] = P_{i//L, i\%L} \qquad (i=0,1,...,L^2-1)

行列$Q$はコスト関数の括弧を展開して出てくる項$p[i]p[j]$の係数を格納していくことで定義できます。具体的には以下の形になります。

\begin{align}
Q[i,j]=
    \begin{cases}
        4    &   \text{($i=j$のとき)}  \\
        -1   &   \text{($p[i] = P_{i//L, i\%L}$と$p[j] = P_{j//L, j\%L}$が隣接しているとき)}
    \end{cases}
\end{align}

ここでは括弧を外した時に出てくる$-2p[i]p[j]$を$Q[i,j]$と$Q[j,i]$に半分ずつ配分しています。こうすると$Q$が対称行列となり、性質が良くなります。これをpythonで書くとこんな感じになります。

Q = np.zeros(shape=(L**2, L**2), dtype=int)
for i in range(L**2):
    yi = i//L
    xi = i%L
    for j in range(L**2):
        yj = j//L
        xj = j%L
        
        if i==j:
            Q[i,j] = 4
        elif is_nearest_neighbor(yi,xi,yj,xj): # 判定関数は適当に作っておく
            Q[i,j] = -1

さて、コスト関数をきれいな形で書き直せたので、ベクトル$p$について最小化すればよいわけですが、すべての成分を自由に動かせるわけではありませんでした。ベクトル$p$のうち$N$個の成分は固定値で動かせないという設定でしたね。こういう時にはベクトル$p$をうまいこと並び替えて、動かせない成分を後のほうに追いやった新たなベクトル$p'$を定義してやります。また

p' =\begin{pmatrix}
p'_1  \\
p'_2 
\end{pmatrix}

と書き、前半の$p'_1$は自由に動かせる値を持った$L^2-N$次元ベクトル、後半の$p'_2$は固定値を持った$N$次元ベクトルであるとします。この並び替えに対応して係数行列$Q$の成分を並び替えた行列$Q'$が定義でき、コスト関数を以下のように書きなおすことができます。

H = p^TQp = p'^TQ'p' = 
\begin{pmatrix}
p'{}_1^T  &  p'{}_2^T 
\end{pmatrix}
\begin{pmatrix}
Q'_{11} & Q'_{12} \\
Q'_{21} & Q'_{22} 
\end{pmatrix}
\begin{pmatrix}
p'{}_1  \\
p'{}_2 
\end{pmatrix}

新しい行列$Q'_{11}, Q'_{12}, Q'_{21}, Q'_{22}$は行列$Q'$を長方形領域に切り分けた各部分で、 $Q'_{21}= Q'{}_{12}^T$が成り立っています。各部分行列の形は$Q'_{11}$が$(L^2-N,L^2-N)$、$Q'_{12}$が$(L^2-N,N)$、 $Q'_{22}$が$(N,N)$です。

こう書いたうえで、$p'_1$を動かしてコスト関数を最小化したいので、平方完成をしていきます。

\begin{align}
& 
\begin{pmatrix}
p'{}_1^T  &  p'{}_2^T 
\end{pmatrix}
\begin{pmatrix}
Q'_{11} & Q'_{12} \\
Q'{}_{12}^T & Q'_{22} 
\end{pmatrix}
\begin{pmatrix}
p'{}_1  \\
p'{}_2 
\end{pmatrix} \\
&= p'{}_1^T Q'_{11}p'{}_1 + p'{}_1^T Q'_{12}p'{}_2 + p'{}_2^T Q'{}_{12}^Tp'{}_1 + p'{}_2^T Q'_{22}p'{}_2 \\
&=(p'{}_1+Q'{}_{11}^{-1}Q'_{12}p'{}_2)^TQ'{}_{11}(p'{}_1+Q'{}_{11}^{-1}Q'_{12}p'{}_2)- p'{}_2^T Q'{}_{12}^{T}Q'{}_{11}^{-1}Q'_{12}p'{}_2 + p'{}_2^T Q'_{22}p'{}_2 
\end{align}

となりますので、

p'{}_1 = -Q'{}_{11}^{-1}Q'_{12}p'{}_2

とすればコストは最小となります。この式が固定の温度$p'{}_2$から自由に動かせる温度$p'{}_1$を決める式です。
また、その時の最小コストは

H(\{p'{}_2\}) = p'{}_2^T\left(   Q'_{22}  - Q'{}_{12}^{T}Q'{}_{11}^{-1}Q'_{12}\right)p'{}_2

であるとわかります。

解法2の計算量

まず上述の行列$Q$を定義する簡単なコードの二重のfor文から明らかなとおり、適当に実装してしまうと行列定義だけで$L^4$の計算量とメモリが必要で、おそらく時間が足りません。加えて、逆行列を含む式の評価には$(L^2)^3=L^6$の時間がかかるためさらに絶望的に思えます。そこで、行列たちの要素がほとんど$0$であること、すなわち疎行列であることに注目し、計算量を下げることが可能です。

疎行列であることを活用すると、「$(i,j)$成分が$Q[i,j]$であること」を非ゼロの成分の数だけ保持しておけばよいため、メモリを大幅に節約できます。また、行列積や逆行列の計算なども専用のアルゴリズムが多数あり、計算量が行列のサイズではなく非ゼロ成分の数でスケールするようになるため、現実的な時間で計算が可能です。またpythonではscipy.sparseからこれらのアルゴリズムを利用可能です。

実装

以下の記述はコンテスト中に書いたコードをもとに作成しており、乱雑なので飛ばしてもOKです

最適な温度$p'{}_1$を知るためには固定温度を並べたベクトル$p'{}_2$、行列たち$Q'{}_{11}, Q'_{12}$が必要ですが、$p'{}_2$に関してはサンプルプログラムを参考に[i * 10 for i in range(N)]のようなものが与えられると仮定します。行列たちの定義には、まず($0~L^2-1$)までを走る元のマスのインデックスを、より小さい行列($Q'{}_{11}, Q'_{12}$など)のインデックスにマップする仕組みが必要です。
固定したい位置を格納したiterableであるfixed_pos (=コンテスト中のlanding_posに対応)が与えらえれているとき、以下のように計算できます。

index_fix = [] # 自由に動かせるマスの番号(0~L^2-1)
index_free= [] # 温度固定のマスの番号(0~L^2-1)
for pos in fixed_pos:
    idx = pos.y*L + pos.x
    index_fix.append(idx)
index_fix_set = set(index_fix)
index_fix = sorted(list(index_fix_set))
for y in range(L):
    for x in range(L):
        i = y*L+x
        if not i in index_fix_set:
            index_free.append(i)   
index_free_set = set(index_free)    

index_fixT  = {v:i for i,v in enumerate(index_fix)} #(0~L^2-1)を(0~L^2-N-1)にマップする転置インデックス
index_freeT = {v:i for i,v in enumerate(index_free)} #(0~L^2-1)を(0~N-1)にマップする転置インデックス

これらのインデックスを使って、疎行列を以下のように定義できます。疎行列の定義にはscipy.sparse のcsr_matrixまたはcsc_matrixを使います。前者は行へのアクセスに優れており、行列積の左側に表れる行列にはこちらを使うとよく、後者は逆、というイメージです。。定義に必要なのは数値、行番号、列番号を表す同じ長さの3つの配列なので、これらを作っていきます。新しく作成する疎行列は元の$L^2$よりも小さいので、先ほどのindex_fixTなどを使ってインデックスを$L^2-N$などに収まるよう変換します。

from scipy.sparse import csr_matrix,csc_matrix

dim_fix = len(index_fix) #=N
# Q11を表す配列。対角成分の4はあらかじめ入れておく。
data11 = [4]*(L**2-dim_fix) 
col11 = list(range(L**2-dim_fix))
row11 = list(range(L**2-dim_fix))

# Q12を表す配列。
data12 = [] 
col12 = []
row12 = []

# Q22を表す配列。対角成分の4はあらかじめ入れておく。
data22 = [4]*(dim_fix) 
col22 = list(range(dim_fix))
row22 = list(range(dim_fix))

# 順番に配列に値を入れていく。もっといい方法がありそうだが、、、
for y in range(L):
    for x in range(L):
        i = y*L+x
        idx = (i+L)%L**2 # (y,x+1)に対応するインデックス
        idy = (i+1)      # (y+1,x)に対応するインデックス
        if idy%L==0: idy-=L
        
        if i in index_free_set and idx in index_free_set:
            data11.append(-1)
            col11.append(index_freeT[i])
            row11.append(index_freeT[idx])
            data11.append(-1)
            col11.append(index_freeT[idx])
            row11.append(index_freeT[i])
        elif i in index_free_set and (idx not in index_free_set):
            data12.append(-1)
            row12.append(index_freeT[i])
            col12.append(index_fixT[idx])
        elif (i not in index_free_set) and idx in index_free_set:
            data12.append(-1)
            row12.append(index_freeT[idx])
            col12.append(index_fixT[i])
        elif (i not in index_free_set) and (idx not in index_free_set):
            data22.append(-1)
            col22.append(index_fixT[i])
            row22.append(index_fixT[idx])
            data22.append(-1)
            col22.append(index_fixT[idx])
            row22.append(index_fixT[i])
            
        if i in index_free_set and idy in index_free_set:
            data11.append(-1)
            col11.append(index_freeT[i])
            row11.append(index_freeT[idy])
            data11.append(-1)
            col11.append(index_freeT[idy])
            row11.append(index_freeT[i])
        elif i in index_free_set and (idy not in index_free_set):
            data12.append(-1)
            row12.append(index_freeT[i])
            col12.append(index_fixT[idy])
        elif (i not in index_free_set) and idy in index_free_set:
            data12.append(-1)
            row12.append(index_freeT[idy])
            col12.append(index_fixT[i])
        elif (i not in index_free_set) and (idy not in index_free_set):
            data22.append(-1)
            col22.append(index_fixT[i])
            row22.append(index_fixT[idy])
            data22.append(-1)
            col22.append(index_fixT[idy])
            row22.append(index_fixT[i])
# SparseEfficiencyWarning:が消えるようにcsrとcscを選択。
Q11 = csc_matrix((data11, (row11, col11) ) , shape=(L**2-dim_fix, L**2-dim_fix))
Q12 = csc_matrix((data12, (row12, col12) ) , shape=(L**2-dim_fix, dim_fix))
Q22 = csr_matrix((data22, (row22, col22) ) , shape=(dim_fix, dim_fix))

Q11.toarray()などで通常のnumpyの密行列表現を得ることができるので、適宜小さい行列で値を確認しながら書いていきます。

ここまでで固定温度$p'{}_2$と行列たち$Q'{}_{11}, Q'{}_{12}, Q'{}_{22}$が得られたので、最適温度$p'{}_1 = -Q'{}_{11}^{-1}Q'_{12}p'{}_2$を計算可能です。一般に行列を用いた数値計算では、精度や計算量の観点から逆行列の計算をしてはいけないといわれています。よって、$p'{}_1 = -Q'{}_{11}^{-1}Q'_{12}p'{}_2$を計算するのではなく、$p'{}_1 \leftarrow \left( Q'{}_{11}p'{}_1 = -Q'_{12}p'{}_2 の解\right)$として計算を行うと高速かつ高精度になります。また、scipy.sparseには疎行列を係数行列として持つ連立一次方程式を高速に解くspsolveという関数が用意されていますのでこれを使うだけです。

from scipy.sparse.linalg import spsolve
p2 = np.array( [i * 10 for i in range(N)] )
p1_opt = spsolve(Q11, -Q12 @ p2)
p1_opt = q1_opt.astype(int)

これですべての温度が決まったので、二次元配列に収めることができます。

temperature=np.zeros(shape(L**2,))
temperature[index_free] = p1_opt
temperature[index_fix]  = p2
temperature = temperature.reshape(self.L,self.L)

以上で、最適な温度配置を求めることができました。

おまけ:さらなる低コストへ

以上の計算で自由に決められる部分の温度$p'{}_1$が最適化され、配置コストはおおよそ(整数化による誤差を除くと)

H(\{p'{}_2\}) = p'{}_2^T\left(   Q'_{22}  - Q'{}_{12}^{T}Q'{}_{11}^{-1}Q'_{12}\right)p'{}_2

まで小さくなります。この値は固定する温度$p'{}_2$の配置で決まっていますので、固定していたつもりのN次元ベクトル$p'{}_2$を動かすことでさらに配置コストを下げることができます。当然$p'{}_2=0$や等温度を選ぶとコストはゼロにできますが、今回のコンテストではのちの処理のために[i * 10 for i in range(N)]のように間隔があいていることが望ましいです。こうした間隔を正則化としてコスト関数に足して最小化してもよいですし、[i * 10 for i in range(N)]の並び替えの範囲で$p'{}_2$を探索する、としてもよいです。

この後者の場合コストの最小化問題は、$0\sim N-1$の置換を$\sigma$として最適な$\sigma$を探す問題

\min H(\{p'{}_2\}) =\min_{\sigma} \sum_{i,j} p'{}_2^T[\sigma(i)]\left(   Q'_{22}  - Q'{}_{12}^{T}Q'{}_{11}^{-1}Q'_{12}\right)[i,j]p'{}_2 [\sigma(j)]

となります。これは二次割当問題という有名な(難しい)問題で、2-optなどのいろいろな近似解法が存在し、ある程度までコストを小さくすることが可能です。

おまけ:Poisson方程式

今回の問題、すなわち配置コスト

H = \sum_{i=0}^{L-1}\sum_{j=0}^{L-1} \left[ ( P_{i,j}- P_{i,(j+1) \% L} )^2+( P_{i,j}- P_{(i+1) \% L, j} )^2 \right]

を適当な固定箇所を除いて最小化する問題は、境界条件付きポアソン方程式を離散近似で解く問題と同値になっており、点電荷を配置したときの定常電場分布や、今回のコンテストのストーリーのように温度固定点をいくつか定めた時の平衡温度分布を求める方程式を表しています。その他の領域でも固体の格子振動や、二次元物質のうねり、画像処理におけるノイズ除去など、様々な場面でPoisson方程式が現れることが知られています。