動機
競技プログラミングの問題を解いていると、これまでに segment tree を使う場面が何回かあったので、実装の仕方を検索してみたのですが、 Fortran で実装しているサンプルが見当たらなかったので、他の言語での実装例を参考にして作ってみました。
segment tree とは
segment tree は 2 つの要素のうち片方を選ぶような操作を、配列の指定した区間の中の全ての要素に対して行うとき、最後まで選ばれる要素を高速に取得することができるデータ構造です。例えば整数の配列に対し、インデックスが $r$ から $l$ までの区間 (以下、 $[r,l]$ と書く) での最小値や最大値を、 minval() や maxval() で愚直に求めるよりも高速に求めることができます。 segment tree は完全二分木でデータを管理しており、葉に元の配列を保管しています。完全二分木の内部ノードにはそのノードを根とする部分木内での操作の結果が保管されており、この結果を利用することで、配列のある区間での最終結果を高速に求めることができます。データ量を $N$ とした場合、愚直な実装では計算量は $O(N)$ ですが、 segment tree を用いると $O(\log N)$ となります。
注意点
この記事で実装している segment tree は整数の配列に対し、 $[l,r]$ の最小値を求めるものとなっております。下のコードの関数 op() を書き換えることで、最大値や部分和などに対応する segment tree にすることができます。ただし、作用させる操作は操作する順番に依存しないことが必要です。 $[l,r]$ が不正 ($l > r$) なクエリの場合は、 segment tree を初期化する際に用いる初期値を返します。また、この記事で実装している query() は再帰的に関数を呼び出す一般的なトップダウンの実装ではなく、ループで回すボトムアップの実装となっております。実装に際してこちらを参考にしました (参考 Segment Tree を少し速くする) 。
他の言語との相違点
他の言語 (例えば C/C++ や Java 、 Python 、 Go など) での実装は 0-indexed で半開区間 $[l,r)$ に対応したものが主流ですが、 Fortran の仕様に合わせて 1-indexed で閉区間 $[l,r]$ に対応させています。
環境
macOS Sierra (10.12.6)
GNU Fortran (GCC) 6.3.0
プログラム
実際に実装したのが以下のものです。
module mod_segment_tree
type t_segment_tree
integer :: deflt
integer :: n, p
integer, pointer :: arr(:) => null()
end type t_segment_tree
contains
integer function op(x,y)
implicit none
integer :: x, y
op = min(x,y)
return
end function op
subroutine init_segtree(st,n,deflt)
implicit none
type(t_segment_tree), intent(inout) :: st
integer, intent(in) :: n
integer, intent(in) :: deflt
integer :: p = 1
do while (p < n)
p = 2*p
end do
st%deflt = deflt
st%n = n
st%p = p
allocate(st%arr(2*p-1))
st%arr = deflt
return
end subroutine init_segtree
subroutine set_default(st,i)
implicit none
type(t_segment_tree), intent(inout) :: st
integer, intent(in) :: i
integer :: x
x = i+st%p-1
st%arr(x) = st%deflt
do while (x > 1)
x = x/2
st%arr(x) = op(st%arr(2*x),st%arr(2*x+1))
end do
return
end subroutine set_default
subroutine update(st,i,v)
implicit none
type(t_segment_tree), intent(inout) :: st
integer, intent(in) :: i
integer, intent(in) :: v
integer :: x
x = i+st%p-1
st%arr(x) = v
do while (x > 1)
x = x/2
st%arr(x) = op(st%arr(2*x),st%arr(2*x+1))
end do
return
end subroutine update
integer function query(st,a,b)
implicit none
type(t_segment_tree), intent(inout) :: st
integer, intent(in) :: a, b
integer :: l, r
query = st%deflt
l = a+st%p-1
r = b+st%p-1
do while (l <= r)
if (mod(l,2) == 1) then
query = op(query,st%arr(l))
l = l+1
end if
if (mod(r,2) == 0) then
query = op(query,st%arr(r))
r = r-1
end if
l = l/2
r = r/2
end do
return
end function query
end module mod_segment_tree
また、実装できているか確認するために以下のプログラムも書きました。
program test_segment_tree
use mod_segment_tree
implicit none
type(t_segment_tree) :: st
integer, parameter :: n = 10
integer :: x(n), i
call random_seed_clock()
call randint(x,10*n)
call init_segtree(st,n,100*n)
write(*,'(a)',advance='no') "array: "
call output(x)
do i = 1, n
call update(st,i,x(i))
end do
write(*,'("query [",i0,",",i0,"]: ",i0)') 1, 10, query(st,1,10)
write(*,'("query [",i0,",",i0,"]: ",i0)') 1, 3, query(st,1,3)
write(*,'("query [",i0,",",i0,"]: ",i0)') 4, 6, query(st,4,6)
write(*,'("query [",i0,",",i0,"]: ",i0)') 7, 9, query(st,7,9)
write(*,'("query [",i0,",",i0,"]: ",i0)') 10, 10, query(st,10,10)
write(*,'("query [",i0,",",i0,"]: ",i0)') 4, 3, query(st,4,3)
stop
contains
subroutine output(a)
implicit none
integer, intent(in) :: a(:)
integer :: n, i
n = size(a)
do i = 1, n
write(*,'(i0,x)',advance='no') a(i)
end do
write(*,*)
return
end subroutine output
subroutine random_seed_clock()
implicit none
integer :: nseed, clock
integer, allocatable :: seed(:)
call system_clock(clock)
call random_seed(size=nseed)
allocate(seed(nseed))
seed = clock
call random_seed(put=seed)
deallocate(seed)
return
end subroutine random_seed_clock
subroutine randint(x,maxint)
implicit none
integer, intent(out) :: x(:)
integer, intent(in) :: maxint
real(8) :: r(size(x))
call random_number(r)
x = int(maxint*r)
return
end subroutine randint
end program test_segment_tree
実行結果は以下のとおりです。
array: 68 58 34 72 51 48 61 6 81 70
query [1,10]: 6
query [1,3]: 34
query [4,6]: 48
query [7,9]: 6
query [10,10]: 70
query [4,3]: 1000
各範囲の最小値が取り出されていることが確認できます。また、 $[4,3]$ のような不正な区間に対するクエリにはデフォルト値が返ってきていることも確認できます。
追記 2019/08/10
不必要な部分を削り、よりオブジェクト指向な書き方に変更しました (ソースコード) 。