AtCoder typical contestでアルゴリズムの勉強 - pythonとfortran

AtCoderにはtypical contestという例題セットがあり、典型的なアルゴリズムの勉強ができます。

講座用というだけあって、解説パワーポイントが問題のページに紹介されており、それを読めばコーディングの指針がわかります。といってもアルゴリズム実装初心者だとそれでもなかなか難しいですが。

https://atc001.contest.atcoder.jp/

こちらをpythonとFortranで実装してみます。

まずPythonで実装し、できるだけ同じアルゴリズムになるようにFortranを実装しました。

# 項目

* [深さ優先探索](https://atc001.contest.atcoder.jp/tasks/dfs_a)
* [Union Find](https://atc001.contest.atcoder.jp/tasks/unionfind_a)
* [高速フーリエ変換](https://atc001.contest.atcoder.jp/tasks/fft_c)

# なぜFortranとPythonか

Fortranだと見てくれる人少ないし、人気言語のPythonで客寄せしてみるか（PythonというかNumpyとFortranが同じくらいのアルゴリズムの表現能力を持っていることを示したい。）

# [深さ優先探索](https://atc001.contest.atcoder.jp/tasks/dfs_a)

迷路の経路存在確認です。最短でもなければ総当りでもないので、深さ優先探索が向いているようです。

ざっくり言うと、チェックするべきマスを未チェック行列に追加し、未チェック行列から１つとってチェックし、隣接するマスを未チェック行列に追加していく処理を繰り返します。
ゴールに着くか、未チェック行列が空になったら終了です。

LIFOのスタックとして未チェック行列を考えると深さ優先探索になるのかなと。
再帰で書くと一部界隈ではイケてるらしいですが、スタックを使います。

```python
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で素直に実装するとこうなります。

```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の直前に

```fortran
   do y = 1, H
      write(*, '(*(g0))')checked(:, y)
   enddo
```
としてやると、H行M列の`T`もしくは`F`の文字列が出力されます。`T`のところが通った道です。


# [Union Find](https://atc001.contest.atcoder.jp/tasks/unionfind_a)

基本的な戦略は解説スライドの通りです。

ということで工夫した点をあげると、

* 木AにノードBとその根Rを接続する時、Rの根をAとするだけでなく、Bも直接Aに繋げます。
* 木Aと木Bを合体させるときは、大きい木に小さい木を繋げます。
* 木の根をチェックする関数を呼び出すたびに、ノードがいる深さが２を超えているかどうか（超えていたらあるノードの直上のノードが根になっていない）を調べ、超えていたらそのノードを根に直接つなげ直す

## Python

```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

```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の様なコンパイラ言語なら行けるかと思い、全てのノードに対してどの木に属しているかのデータを持たせ、クエリの度に更新させてみましたが、ちょっと無理でした。


# [高速フーリエ変換](https://atc001.contest.atcoder.jp/tasks/fft_c)

問題がフーリエ変換と同値であることは、問題ページにあるスライドが詳しいですので、そちらをご覧ください。

pythonではnumpyにfftモジュールがあるので使ってしまいます。AtCoderではnumpyとscipyが使えるのです。

```python
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 - 高速フーリエ変換](https://ja.wikipedia.org/wiki/%E9%AB%98%E9%80%9F%E3%83%95%E3%83%BC%E3%83%AA%E3%82%A8%E5%A4%89%E6%8F%9B)の内容のほうが分かりやすかったので、こちらをもとに実装します。

まず、基本的な離散フーリエ変換は、長さ$N$の配列 $\vec{x}$ に対して $W = \exp(-2\pi i/N)$ として

```math
\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$に対して 

```math
X_k = \sum^{N-1}_{j=0} W^{jk}x_j
```
これを行列表記で

```math
\vec{X} = W\vec{x}
```

となります。これをFortranのコードにしてみます。

```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$とします。

```math
X_n = \sum^{N-1}_{k=0} W^{nk}x_k
```
に対し、$N=PQ$のとき、

```math
0\leq s,p \leq P -1 \\
0\leq r,q \leq Q-1
```
なる $p, q, r, s$を用いて、

```math
n=sQ +r\\
k = qP + p
```
と添字を書きなおしてやると、

```math
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)$なので

```math
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$のベクトルのフーリエ変換の変換行列の要素とみなせます。

こう考えると、添字を書きなおしたあとに出てきた

```math
 \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$を選んだかで生じます。
すると

```math
X_{sQ+r} = \sum^{P-1}_{p=0}W^{spQ}\xi_r^p
```
となります。$\xi$の添字がこれまでのフーリエ変換とは上下逆になっただけで、これは$P$の長さを持つ配列$\xi^p$をフーリエ変換したというように見えなくもないです。

1. $N=PQ$とする
2. $\vec{x}$から$P$で割って$p$余る番目の要素$Q$個を取り出し、フーリエ変換
3. 2を$P$回繰り返し
4. $P$回繰り返したものを並べ、係数を掛けて、長さ$P$のフーリエ変換ぽい感じで$\vec{X}$を構築

こんなふうに因数分解した因数でフーリエ変換を分割するアルゴリズムを利用します。

## ２の累乗の場合

ここで$P=2$、$Q=N/2$とします。$\sum_{p=0}^{P-1}$はどうせ２項だけなので書き下してしまうと、

```math
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$に書き換えつつ、

```math
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$のベクトルのフーリエ変換（２分割版）は、

1. $x_n$の奇数番目の成分 $x^{odd}\_k$と偶数番目の成分$x^{even}\_k$に分け、それぞれでフーリエ変換 $x^{odd}\_k → X^{odd}\_n $と $x^{even}\_k → X^{even}\_n $を行う
2. $W^nX^{odd}\_n$と$X^{even}\_n$を$N$次元のベクトルとして一列に並べ、前半・後半をそれぞれ１つの成分と思い、２成分のフーリエ変換をする
3. $X\_n$のできあがり

という感じです。で、1.の$x^{odd}\_k → X^{odd}\_n $と $x^{even}\_k → X^{even}\_n $のフーリエ変換もまた分割フーリエ変換を行い、再帰的にさらに短いベクトルでのフーリエ変換に分割していきます。

そして十分短くなった時、例えば $N = 2$の時は$W = \exp(-\pi i) = -1 $なので

```math
\vec{X} = \left(\begin{array}{cc}
1 & 1 \\
1 & -1 \end{array}\right)\vec{x}
```
$N = 4$の時は$W = \exp(-\pi i/2) = -i $なので

```math
\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になったらこれらの行列を使って変換してやって、元のベクトルのフーリエ変換したベクトルを組み上げていく形になります。

## 逆変換

```math
\vec{x} = \frac{1}{N}W^{-1}\vec{X}
```
ですが、$W$の定義から

```math
\vec{x} = \frac{1}{N}\overline{W\overline{\vec{X}}}
```
と、ベクトル成分の複素数を共役にしたものをフーリエ変換し、それをまた共役を取ればいいので、改めて実装する必要はありません。


## 以上をまとめたFortranコード

```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$以外の長さのものが渡されることは考えていません。

# 終わりに

ライブラリ使いましょう。