AtCoderにはtypical contestという例題セットがあり、典型的なアルゴリズムの勉強ができます。
講座用というだけあって、解説パワーポイントが問題のページに紹介されており、それを読めばコーディングの指針がわかります。といってもアルゴリズム実装初心者だとそれでもなかなか難しいですが。
こちらをpythonとFortranで実装してみます。
まずPythonで実装し、できるだけ同じアルゴリズムになるようにFortranを実装しました。
項目
なぜFortranとPythonか
Fortranだと見てくれる人少ないし、人気言語のPythonで客寄せしてみるか（PythonというかNumpyとFortranが同じくらいのアルゴリズムの表現能力を持っていることを示したい。）
深さ優先探索
迷路の経路存在確認です。最短でもなければ総当りでもないので、深さ優先探索が向いているようです。
ざっくり言うと、チェックするべきマスを未チェック行列に追加し、未チェック行列から１つとってチェックし、隣接するマスを未チェック行列に追加していく処理を繰り返します。
ゴールに着くか、未チェック行列が空になったら終了です。
LIFOのスタックとして未チェック行列を考えると深さ優先探索になるのかなと。
再帰で書くと一部界隈ではイケてるらしいですが、スタックを使います。
H, W = map(int, input().split())
M = [input() for _ in range(H)]
start_idx = "".join(M).index("s")
next_list = [(start_idx % W, start_idx // W),]
checked = [[False]*W for _ in range(H)]
def touch(xj, yj):
if xj < 0 or yj < 0 or xj > W - 1 or yj > H - 1:
return
if M[yj][xj] == "#" or checked[yj][xj]:
return
next_list.append((xj, yj))
while len(next_list) > 0:
xi, yi = next_list.pop()
checked[yi][xi] = True
if M[yi][xi] == "g":
print("Yes")
exit()
touch(xi - 1, yi)
touch(xi, yi - 1)
touch(xi + 1, yi)
touch(xi, yi + 1)
print("No")
これをFortranで素直に実装するとこうなります。
program main
implicit none
integer :: H, W, sx, sy(1), remain, x, y
character(len=500), allocatable :: M(:)
integer, allocatable :: uncheck(:, :), s(:)
logical, allocatable :: checked(:, :)
logical :: found
read(*, *) H, W
allocate(M(H))
allocate(uncheck(2, H*W))
allocate(checked(W, H))
read(*, *) M
s = index(M, "s")
sx = maxval(s)
sy = maxloc(s)
uncheck(:, 1) = [sx, sy]
remain = 1
checked = .False.
found = .False.
do while(remain > 0)
x = uncheck(1, remain)
y = uncheck(2, remain)
remain = remain - 1
checked(x, y) = .True.
if(M(y)(x:x) == "g")then
found = .True.
exit
endif
call touch(x - 1, y)
call touch(x + 1, y)
call touch(x, y - 1)
call touch(x, y + 1)
enddo
if (found) then
write(*, *) "Yes"
else
write(*, *) "No"
endif
deallocate(M)
deallocate(uncheck)
deallocate(checked)
stop
contains
subroutine touch(ix, iy)
integer, intent(in) :: ix, iy
if(ix < 1 .or. iy < 1 .or. ix > W .or. iy > H) return
if( M(iy)(ix:ix) == "#" .or. checked(ix, iy)) return
remain = remain + 1
uncheck(:, remain) = [ix, iy]
end subroutine touch
end program main
変数宣言とメモリ関連の分長くなりましたが、基本的には素直に翻訳できたかと思います。
なお、実際マップ上のどこを見たかをチェックするには、deallocateの直前に
do y = 1, H
write(*, '(*(g0))')checked(:, y)
enddo
としてやると、H行M列のTもしくはFの文字列が出力されます。Tのところが通った道です。
Union Find
基本的な戦略は解説スライドの通りです。
ということで工夫した点をあげると、
- 木AにノードBとその根Rを接続する時、Rの根をAとするだけでなく、Bも直接Aに繋げます。
- 木Aと木Bを合体させるときは、大きい木に小さい木を繋げます。
- 木の根をチェックする関数を呼び出すたびに、ノードがいる深さが２を超えているかどうか（超えていたらあるノードの直上のノードが根になっていない）を調べ、超えていたらそのノードを根に直接つなげ直す
Python
N, Q = map(int, input().split())
follow = [-1]*N # -1 means root
num_follower = [1]*N
def carry_up(A, root):
pA = A
while pA != root:
pA, follow[pA] = follow[pA], root
def root_idx_of(A):
r = A
while follow[r] > -1:
r = follow[r]
if r != follow[A]:
carry_up(A, r)
return r
def connected(A, B):
return root_idx_of(A) == root_idx_of(B)
def connect(A, B):
rA = root_idx_of(A)
rB = root_idx_of(B)
if rA == rB:
return
if num_follower[rA] < num_follower[rB]:
follow[rA] = rB
follow[A] = rB
num_follower[rB] += num_follower[rA]
else:
follow[rB] = rA
follow[B] = rA
num_follower[rA] += num_follower[rB]
for _ in range(Q):
P, A, B = map(int, input().split())
if P == 0:
connect(A, B)
elif P == 1:
if connected(A, B):
print("Yes")
else:
print("No")
Fortran
program main
implicit none
integer :: N, query(3), Q, i, A, B
integer, allocatable :: root(:), counter(:)
read(*, *) N, Q
root = [(0, i=1,N)] ! root(i) == -1 means i is root
counter = [(1, i=1, N)]
do i = 1, Q
read(*, *) query
A = query(2) + 1
B = query(3) + 1
if(query(1) == 0)then
call connect(A, B)
endif
if(query(1) == 1)then
if (root_idx(A) == root_idx(B))then
write(*, *) "Yes"
else
write(*, *) "No"
endif
endif
enddo
deallocate(root)
contains
subroutine carry_up(x, r)
integer, intent(in) :: x, r
integer :: rx, trx
rx = x
do while(rx /= r)
trx = root(rx)
root(rx) = r
rx = trx
enddo
end subroutine carry_up
function root_idx(x) result(r)
integer, intent(in) :: x
integer :: r
r = x
do while(root(r) /= 0)
r = root(r)
enddo
if (r /= root(x)) call carry_up(x, r)
end function root_idx
subroutine connect(n1, n2)
integer, intent(in) :: n1, n2
integer :: r1, r2
r1 = root_idx(n1)
r2 = root_idx(n2)
if(r1 == r2) return
if(counter(r1) > counter(r2))then
root(n2) = r1
root(r2) = r1
counter(r1) = counter(r1) + counter(r2)
else
root(n1) = r2
root(r1) = r2
counter(r2) = counter(r2) + counter(r1)
endif
end subroutine connect
end program main
ちなみに、Fortranの様なコンパイラ言語なら行けるかと思い、全てのノードに対してどの木に属しているかのデータを持たせ、クエリの度に更新させてみましたが、ちょっと無理でした。
高速フーリエ変換
問題がフーリエ変換と同値であることは、問題ページにあるスライドが詳しいですので、そちらをご覧ください。
pythonではnumpyにfftモジュールがあるので使ってしまいます。AtCoderではnumpyとscipyが使えるのです。
import numpy as np
N = int(input())
n0 = 2**int(np.ceil(np.log2(2*N - 1)))
A = np.zeros(n0)
B = np.zeros(n0)
for i in range(N):
A[i], B[i] = map(int, input().split())
C = np.fft.ifft( np.fft.fft(A)*np.fft.fft(B) )
ret = ["0"] + ["{}".format(int(ci)) for ci in np.real(C[:2*N - 1] + 0.5)]
print("\n".join(ret))
C = np.fft.ifft(以下略 の一行で、ベクトルAとBをフーリエ変換しつつ要素積をとり、逆変換してCに戻しています。この問題の肝の部分です。
n0はN以上となる最小の2のべき乗数を求めています。
これをこれからFortranで書いていきます。
が、FortranでFFTをしようとするとAtCoderではFFTWのリンクが出来ないので、自分で用意してやる必要があります。スライドでも高速フーリエ変換関数を自分で書くことを想定しているようですし。
基本のフーリエ変換
ここからしばらく数式ばかりでてきます。「アルゴリズムはプログラムのコードで勉強できます」という人は最後の「以上をまとめたFortranコード」まで飛ばしてください。
個人的にフーリエ変換についてはAtCoderの解説よりもWikipedia - 高速フーリエ変換の内容のほうが分かりやすかったので、こちらをもとに実装します。
まず、基本的な離散フーリエ変換は、長さ$N$の配列 $\vec{x}$ に対して $W = \exp(-2\pi i/N)$ として
\vec{X} = \left(\begin{array}{ccccccc}
W^0 & W^0 & W^0 & \cdots & W^0 & & W^0\\
W^0 & W^1 & W^2 & & W^k & & W^{N-1}\\
W^0 & W^2 & W^4 & & & &\\
\vdots & & & \ddots & & &\\
W^0 & W^j & & & W^{jk} & & \\
& & & & & \ddots & \\
W^0 & W^{N-1}& & \cdots & & & W^{(N-1)(N-1)}
\end{array}\right)\vec{x}
あるいは $k = 0, 1, \cdots, N-1$に対して
X_k = \sum^{N-1}_{j=0} W^{jk}x_j
これを行列表記で
\vec{X} = W\vec{x}
となります。これをFortranのコードにしてみます。
double precision, parameter :: PI = acos(-1.d0)
function dft_dft(f, N) result(ff)
complex(kind(0d0)), intent(in) :: f(N)
integer, intent(in) :: N
complex(kind(0d0)) :: ff(N), W(N, N)
complex(kind(0d0)) :: Wn
integer :: i, j
Wn = exp((0.d0, -1.d0)*2.d0*PI/N)
forall(i = 1:N, j = 1:N) W(i, j) = Wn**((i - 1)*(j - 1))
ff = matmul(W, f)
end function dft_dft
これはベクトルの長さは１以上であれば何でも大丈夫です。
ただしこのままだと計算時間、メモリー消費量ともに $O(N^2)$ で、 $N$が大きいテストケースでRuntime Errorとなりました（なぜTLEとかMLEじゃないんだろう）。
なので高速フーリエ変換を実装します。
高速フーリエ変換
ここは特にわかりにくい書き方になってしまったかなと。
一応僕なりの解説を書いておきますが、正直これまでライブラリの利用で満足していて勉強してこなかった人間のまとめですので、ご容赦ください。
Wikiの高速フーリエ変換のページの「原理の説明」の節をそのまま使います。ただ、Wikipediaでは$f(n)$, $F(n)$となっているのを、この記事のこれまでのに合わせて$x_n$、$X_n$とします。
X_n = \sum^{N-1}_{k=0} W^{nk}x_k
に対し、$N=PQ$のとき、
0\leq s,p \leq P -1 \\
0\leq r,q \leq Q-1
なる $p, q, r, s$を用いて、
n=sQ +r\\
k = qP + p
と添字を書きなおしてやると、
X_n = X_{sQ+r} = \sum^{P-1}_{p=0}W^{(sQ+r)(qP+p)}x_{qP+p}\\
= \sum^{P-1}_{p=0}W^{spQ}\left( W^{pr} \sum^{Q-1}_{q=0}W^{rqP}x_{qP+p}\right)
ここで、$W^{spQ}$や$W^{rqP}$といった因子について考えてみます。
長さが$P$のベクトル$y_p$のフーリエ変換を考えてみます。$W = \exp\left(-\frac{2\pi i}{N}\right)$なので
Y_s = \sum^{P-1}_{p=0}\exp\left(-sp\frac{2\pi i}{P}\right)y_p = \sum^{P-1}_{p=0} W^{spQ}y_p
となり、$W^{spQ}$は長さ$P$のベクトルのフーリエ変換の変換行列の要素とみなせます。
こう考えると、添字を書きなおしたあとに出てきた
\sum^{Q-1}_{q=0}W^{rqP}x_{qP+p}
は、 $P$で割るとあまりが$p$になる$k$を添え字に持つ$x_k$を並べた長さ$Q$のベクトル のフーリエ変換に相当します。
それに$W^{pr}$を掛けたベクトルを$\xi_r^p=W^{pr}\sum^{Q-1}_{q=0}W^{rqP}x_{qP+p}$とします。ややこしいですが$\xi_r^p$は長さ$Q$のベクトルで、それが$P$本ありますので、添字が２つあります。$r$がベクトルの次元を示し、$p$がベクトルの番号を示します。ベクトルの種類は、$P$で割るとあまりが$p$になるという時に、どの$p$を選んだかで生じます。
すると
X_{sQ+r} = \sum^{P-1}_{p=0}W^{spQ}\xi_r^p
となります。$\xi$の添字がこれまでのフーリエ変換とは上下逆になっただけで、これは$P$の長さを持つ配列$\xi^p$をフーリエ変換したというように見えなくもないです。
- $N=PQ$とする
- $\vec{x}$から$P$で割って$p$余る番目の要素$Q$個を取り出し、フーリエ変換
- 2を$P$回繰り返し
- $P$回繰り返したものを並べ、係数を掛けて、長さ$P$のフーリエ変換ぽい感じで$\vec{X}$を構築
こんなふうに因数分解した因数でフーリエ変換を分割するアルゴリズムを利用します。
２の累乗の場合
ここで$P=2$、$Q=N/2$とします。$\sum_{p=0}^{P-1}$はどうせ２項だけなので書き下してしまうと、
X_n = X_{sQ+r}
= \left( \sum^{Q-1}_{q=0}W^{2qr}x_{2q}\right) + W^{sQ}\left( W^{r} \sum^{Q-1}_{q=0}W^{2qr}x_{2q+1}\right)
ですが、$s$が0または1、$W^Q = \exp(-i\pi) = -1$と言うことを考えると、$r$を$n$に書き換えつつ、
X_r = \left\{
\begin{array}{ll}
\left( \sum^{Q-1}_{q=0}W^{2qn}x_{2q}\right) + \left( W^{n} \sum^{Q-1}_{q=0}W^{2qn}x_{2q+1}\right) (n < Q)\\
\left( \sum^{Q-1}_{q=0}W^{2qn}x_{2q}\right) - \left( W^{n - Q} \sum^{Q-1}_{q=0}W^{2qn}x_{2q+1}\right) (n \geq Q)
\end{array}
\right.
とできます。
そして$\left(\quad\right)$内は先でも述べたとおり長さ$Q$のフーリエ変換とみなせます。
つまり、長さ$N$のベクトルのフーリエ変換（２分割版）は、
- $x_n$の奇数番目の成分 $x^{odd}_k$と偶数番目の成分$x^{even}_k$に分け、それぞれでフーリエ変換 $x^{odd}_k → X^{odd}_n $と $x^{even}_k → X^{even}_n $を行う
- $W^nX^{odd}_n$と$X^{even}_n$を$N$次元のベクトルとして一列に並べ、前半・後半をそれぞれ１つの成分と思い、２成分のフーリエ変換をする
- $X_n$のできあがり
という感じです。で、1.の$x^{odd}_k → X^{odd}_n $と $x^{even}_k → X^{even}_n $のフーリエ変換もまた分割フーリエ変換を行い、再帰的にさらに短いベクトルでのフーリエ変換に分割していきます。
そして十分短くなった時、例えば $N = 2$の時は$W = \exp(-\pi i) = -1 $なので
\vec{X} = \left(\begin{array}{cc}
1 & 1 \\
1 & -1 \end{array}\right)\vec{x}
$N = 4$の時は$W = \exp(-\pi i/2) = -i $なので
\vec{X} = \left(\begin{array}{cccc}
1 & 1 & 1 & 1\\
1 & -i & -1 & i \\
1 & -1 & 1 & -1 \\
1 & i & -1 & -i
\end{array}\right)\vec{x}
となります。
$2^N$のベクトルをどんどん半分にしていき、2または4になったらこれらの行列を使って変換してやって、元のベクトルのフーリエ変換したベクトルを組み上げていく形になります。
逆変換
\vec{x} = \frac{1}{N}W^{-1}\vec{X}
ですが、$W$の定義から
\vec{x} = \frac{1}{N}\overline{W\overline{\vec{X}}}
と、ベクトル成分の複素数を共役にしたものをフーリエ変換し、それをまた共役を取ればいいので、改めて実装する必要はありません。
以上をまとめたFortranコード
module fft
implicit none
double precision, parameter :: PI = acos(-1.d0)
contains
function dft_atom(f, N) result(ff)
complex(kind(0d0)), intent(in) :: f(N)
integer, intent(in) :: N
complex(kind(0d0)) :: ff(N)
if ( N == 2 ) then
ff(1) = f(1) + f(2)
ff(2) = f(1) - f(2)
elseif( N == 4 )then
ff(1) = sum(f)
ff(2) = f(1) - (0.d0, 1.d0)*f(2) - f(3) + (0.d0, 1.d0)*f(4)
ff(3) = f(1) - f(2) + f(3) - f(4)
ff(4) = f(1) + (0.d0, 1.d0)*f(2) - f(3) - (0.d0, 1.d0)*f(4)
else
ff = 0.d0
endif
end function dft_atom
recursive function dft_fft(f, N) result(ff)
complex(kind(0d0)), intent(in) :: f(N)
integer, intent(in) :: N
integer :: M, i
complex(kind(0d0)) :: ff(N), Wn
complex(kind(0d0)), allocatable :: fodd(:), feven(:), W(:)
if ( N < 5) then
ff = dft_atom(f, N)
else
M = N/2
Wn = exp( -2*PI*(0.d0, 1.d0)/N)
W = [(Wn**i, i=0,M-1)]
fodd = W*dft_fft(f(2:N:2), M)
feven = dft_fft(f(1:N:2), M)
ff = [(feven + fodd), (feven - fodd)]
endif
end function dft_fft
function dft_ifft(f, N) result(ff)
complex(kind(0d0)), intent(in) :: f(N)
integer, intent(in) :: N
complex(kind(0d0)) :: ff(N)
ff = conjg(dft_fft(conjg(f),N))/N
end function dft_ifft
end module fft
program main
use fft
implicit none
integer :: N, i, N2
integer, ALLOCATABLE :: A(:), B(:), ret(:)
complex(kind(0d0)), allocatable :: fA(:), fB(:)
read(*, *) N
N2 = 2**ceiling(log(2*N - 1.d0)/log(2.d0), kind=1)
allocate(A(N2))
allocate(B(N2))
A = 0
B = 0
do i = 1, N
read(*, *) A(i), B(i)
enddo
fA = dft_fft(A*(1.d0, 0.d0), N2)
fB = dft_fft(B*(1.d0, 0.d0), N2)
fA = dft_ifft(fA*fB, N2)
ret = [0, nint(real(fA))]
write(*, '(*(i0, :, /))') ret(:2*N)
stop
end program main
Pythonで答えのベクトルCに当たるものを、今回はfAで使いまわしています。
頑張ればPythonみたく１行で変換と積と逆変換ができたかもしれない。
なお、$2^N$以外の長さのものが渡されることは考えていません。
終わりに
ライブラリ使いましょう。