Fortran で binary indexed tree を作ってみた

動機

AtCoder の問題を解いているとしばしば binary indexed tree (BIT, Fenwick tree) が登場するので Fortran で実装してみました。 binary indexed tree の実装は 1-indexed なので Fortran だと素直に実装できます。また、今回は普段よりオブジェクト指向を意識して実装しました。

binary indexed tree とは

binary indexed tree は配列の部分和を高速に求めることができる木構造です。 Peter M. Fenwick 氏が 1994 年に提案したため、しばしば Fenwick tree とも呼ばれます。また、部分和がある数以上になる最小のインデックスを高速に求めることもできます。データ数を $N$ とするとき、データの更新、部分和の計算、部分和がある数以上になる最小のインデックスの探索がどれも $O(\log N)$ で実現できます。 binary indexed tree の用途としては、部分和の計算や部分和がある数以上になる最小のインデックスの探索はもちろん、 LIS (longest increasing subsequence) の長さや転倒数 (その配列が昇順にソートされている場合から何回要素の入れ替えが生じているか) の計算にも用いられます。

環境

macOS Sierra (10.12.6)
GNU Fortran (GCC) 6.3.0

プログラム

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

fenwick_tree.f08

module mod_fenwick_tree

  type t_fenwick_tree
    private
    integer :: n, p
    integer, pointer :: arr(:) => null()
  contains
    procedure :: init => init_bit
    procedure :: release => release_bit
    procedure :: add => get_addition
    procedure :: sum => get_summation
    procedure :: range => get_partial_summation
    procedure :: val => get_value
    procedure :: lower_bound => get_lower_bound
    procedure :: update => get_updated
    procedure :: max => get_maximum
  end type t_fenwick_tree

contains

  ! BIT の初期化
  subroutine init_bit(this,n)
    implicit none
    class(t_fenwick_tree), intent(inout) :: this
    integer, intent(in) :: n
    integer :: p = 1

    this%n = n
    allocate(this%arr(n))
    this%arr = 0
    do while (lshift(p,1) <= n)
      p = lshift(p,1)
    end do
    this%p = p
    return
  end subroutine init_bit

  ! BIT の解放
  subroutine release_bit(this)
    implicit none
    class(t_fenwick_tree), intent(inout) :: this

    if (associated(this%arr)) deallocate(this%arr)
    return
  end subroutine release_bit

  ! BIT の i 番目の要素に v を加算
  subroutine get_addition(this,i,v)
    implicit none
    class(t_fenwick_tree), intent(inout) :: this
    integer, intent(in) :: i
    integer, intent(in) :: v
    integer :: x

    x = i
    do while (x <= this%n)
      this%arr(x) = this%arr(x)+v
      x = x+and(x,-x)
    end do
    return
  end subroutine get_addition

  ! BIT の i 番目の要素までの総和 (i 番目含む)
  function get_summation(this,i) result(s)
    implicit none
    class(t_fenwick_tree), intent(in) :: this
    integer, intent(in) :: i
    integer :: x
    integer :: s

    s = -1
    if (i > this%n) return
    s = 0
    x = i
    do while (x > 0)
      s = s+this%arr(x)
      x = x-and(x,-x)
    end do
    return
  end function get_summation

  ! BIT の [l,r] の部分和 (両端含む)
  function get_partial_summation(this,l,r) result(s)
    implicit none
    class(t_fenwick_tree), intent(in) :: this
    integer, intent(in) :: l, r
    integer :: i, j
    integer :: s

    s = -1
    if (l > r) return
    s = 0
    i = l-1
    j = r
    do while (j > i)
      s = s+this%arr(j)
      j = j-and(j,-j)
    end do
    do while (i > j)
      s = s-this%arr(i)
      i = i-and(i,-i)
    end do
    return
  end function get_partial_summation

  ! BIT の i番目の要素の取得
  function get_value(this,i) result(v)
    implicit none
    class(t_fenwick_tree), intent(in) :: this
    integer, intent(in) :: i
    integer :: x
    integer :: v

    v = get_partial_summation(this,i,i)
    return
  end function get_value

  ! bit%sum(k) >= v を満たす最小の k
  function get_lower_bound(this,v) result(x)
    implicit none
    class(t_fenwick_tree), intent(in) :: this
    integer, intent(in) :: v
    integer :: x, k
    integer :: w

    x = 0
    if (v <= 0) return
    k = this%p
    w = v
    do while (k > 0)
      if (x+k <= this%n .and. this%arr(x+k) < w) then
        w = w-this%arr(x+k)
        x = x+k
      end if
      k = rshift(k,1)
    end do
    x = x+1
    return
  end function get_lower_bound

  ! LIS の長さを求める際の、 i 番目までの最大値の更新
  subroutine get_updated(this,i,v)
    implicit none
    class(t_fenwick_tree), intent(inout) :: this
    integer, intent(in) :: i
    integer, intent(in) :: v
    integer :: x

    x = i
    do while (x <= this%n)
      if (this%arr(x) < v) this%arr(x) = v
      x = x+and(x,-x)
    end do
    return
  end subroutine get_updated

  ! LIS の長さを求める際の、 i 番目までの最大値の取得
  function get_maximum(this,i) result(m)
    implicit none
    class(t_fenwick_tree), intent(in) :: this
    integer, intent(in) :: i
    integer :: x, m

    x = i
    m = 0
    do while (x > 0)
      m = max(m,this%arr(x))
      x = x-and(x,-x)
    end do
    return
  end function get_maximum

  ! LIS の長さ (forall i a(i) <= len(a) を仮定)
  function len_of_lis(a) result(m)
    implicit none
    integer, intent(in) :: a(:)
    type(t_fenwick_tree) :: bit
    integer :: m, n, i, t

    n = size(a)
    call init_bit(bit,n)
    m = 0
    do i = 1, n
      t = get_maximum(bit,a(i)-1)+1
      m = max(m,t)
      call get_updated(bit,a(i),t)
    end do
    return
  end function len_of_lis

  ! 転倒数 (forall i a(i) <= len(a) を仮定)
  function inv_num(a) result(m)
    implicit none
    integer, intent(in) :: a(:)
    type(t_fenwick_tree) :: bit
    integer :: m, n, i, t

    n = size(a)
    call init_bit(bit,n)
    m = 0
    do i = 1, n
      call get_addition(bit,a(i),1)
      m = m+i-get_summation(bit,a(i))
    end do
    return
  end function inv_num

end module mod_fenwick_tree

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

test_fenwick_tree.f08

program test_fenwick_tree
  use mod_fenwick_tree
  implicit none
  type(t_fenwick_tree) :: bit
  integer, parameter :: n = 7
  integer :: x(n), y(n), i

  do i = 1, n
    x(i) = i*10**(i-1)
  end do

  call bit%init(n)

  write(*,'(a)',advance='no') "original array: "
  call output(x)

  do i = 1, n
    call bit%add(i,x(i))
  end do

  do i = 1, n
    y(i) = bit%val(i)
  end do

  write(*,'(a)',advance='no') "added array: "
  call output(y)

  do i = 1, n
    write(*,'("sum [:",i0,"]: ",i0)') i, bit%sum(i)
  end do
  write(*,'("sum [:",i0,"]: ",i0)') n+1, bit%sum(n+1)

  write(*,'("range [",i0,":",i0,"]: ",i0)') 1, 2, bit%range(1,2)
  write(*,'("range [",i0,":",i0,"]: ",i0)') 3, 5, bit%range(3,5)
  write(*,'("range [",i0,":",i0,"]: ",i0)') 6, 7, bit%range(6,7)
  write(*,'("range [",i0,":",i0,"]: ",i0)') 1, 7, bit%range(1,7)
  write(*,'("range [",i0,":",i0,"]: ",i0)') 2, 6, bit%range(2,6)
  write(*,'("range [",i0,":",i0,"]: ",i0)') 5, 4, bit%range(5,4)

  write(*,'("lower bound of ",i0,": ",i0)') 300, bit%lower_bound(300)
  write(*,'("lower bound of ",i0,": ",i0)') 321, bit%lower_bound(321)
  write(*,'("lower bound of ",i0,": ",i0)') 322, bit%lower_bound(322)
  write(*,'("lower bound of ",i0,": ",i0)') 7654321, bit%lower_bound(7654321)
  write(*,'("lower bound of ",i0,": ",i0)') 7654322, bit%lower_bound(7654322)
  write(*,'("lower bound of ",i0,": ",i0)') 10000000, bit%lower_bound(10000000)
  write(*,'("lower bound of ",i0,": ",i0)') 0, bit%lower_bound(0)
  write(*,'("lower bound of ",i0,": ",i0)') -30, bit%lower_bound(-30)

  x = (/1,2,3,4,5,6,7/)
  write(*,'(a)',advance='no') "array: "
  call output(x)
  write(*,'("length of lis: ",i0)') len_of_lis(x)
  write(*,'("value of inversion number: ",i0)') inv_num(x)

  x = (/2,1,2,3,4,3,2/)
  write(*,'(a)',advance='no') "array: "
  call output(x)
  write(*,'("length of lis: ",i0)') len_of_lis(x)
  write(*,'("value of inversion number: ",i0)') inv_num(x)

  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

end program test_fenwick_tree

実行結果は以下のとおりです。

original array: 1 20 300 4000 50000 600000 7000000
added array: 1 20 300 4000 50000 600000 7000000
sum [:1]: 1
sum [:2]: 21
sum [:3]: 321
sum [:4]: 4321
sum [:5]: 54321
sum [:6]: 654321
sum [:7]: 7654321
sum [:8]: -1
range [1:2]: 21
range [3:5]: 54300
range [6:7]: 7600000
range [1:7]: 7654321
range [2:6]: 654320
range [5:4]: -1
lower bound of 300: 3
lower bound of 321: 3
lower bound of 322: 4
lower bound of 7654321: 7
lower bound of 7654322: 8
lower bound of 10000000: 8
lower bound of 0: 0
lower bound of -30: 0
array: 1 2 3 4 5 6 7
length of lis: 7
value of inversion number: 0
array: 2 1 2 3 4 3 2
length of lis: 4
value of inversion number: 5

これで区間和や LIS や転倒数を高速に求められるようになりました。

参考

フェニック木 (Wikipedia)
k 番目の値を高速に取り出せるデータ構造のまとめ - BIT上二分探索や平衡二分探索木など
Binary Indexed Tree のはなし - hos.ac
fortranでもオブジェクト指向したい！

追記 2019/08/10

コメントを受けて少し書き直しました (ソースコード) 。