Fortran で segment tree を作ってみた

動機

競技プログラミングの問題を解いていると、これまでに 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

プログラム

実際に実装したのが以下のものです。

segment_tree.f08

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

また、実装できているか確認するために以下のプログラムも書きました。

test_segment_tree.f08

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

不必要な部分を削り、よりオブジェクト指向な書き方に変更しました (ソースコード) 。