[LT発表会]Ruby meets Go

こんにちは！

今回は前回のLT発表の内容をご紹介いたします。

olivaでは毎月の帰社日で全員集まった際に、LT発表として技術系の発表を行なっています。

前回私の発表した内容をお伝えいたします。

概要

競技プログラミングの問題を題材に、RubyとGoの簡単な速度比較
RubyのコードにGoを組み込んで処理速度アップを図ってみる

Rubyは本当に遅いのか？

今回の題材はこちらの問題↓

AtCoder Regular Contest 104 B – DNA Sequence

ポイントは、競技プログラミング界においては、「10の7乗はおそらく間に合う」レベルの計算量となります。

書いたコード１

まず以下のように解いてみました。（実際にAtCoder参加中に急いで書いたものなので、コードについてはご愛敬で）

（奇数個では相補鎖は作れないので、偶数個だけ判定）

2個の分割から全てまでの部分列を取得
Aの個数==Tの個数, Cの個数==Gの個数　をカウント

n,s=gets.split
n=n.to_i
s=s.chomp.split('')
ans = 0
(1..n/2+1).each do |i|
  next if i*2>n
  s.each_cons(i*2) do |si|
    ans+= 1 if si.count("A")==si.count("T") && si.count("C")==si.count("G")
  end
end
p ans

n,s=gets.split

n=n.to_i

s=s.chomp.split('')

ans = 0

(1..n/2+1).each do |i|

next if i*2>n

s.each_cons(i*2) do |si|

ans+= 1 if si.count("A")==si.count("T") && si.count("C")==si.count("G")

end

p ans

間に合いませんでした・・・。

カウント部分で、再度全部チェックしているため、そこが遅いと考えて修正しました。

書いたコード２

（奇数個では相補鎖は作れないので、偶数個だけ判定）

2個の分割から全てまでの部分列を取得
A個数==Tの個数, Cの個数==Gの個数　をカウント

カウントのループを1回でできるように改善

n=n.to_i
s=s.chomp.split('')
ans = 0
(0...n).each do |i|
  a,t,c,g = 0,0,0,0
  (i...n).each do |j|
    case s[j]
    when "A"
      a += 1
    when "T"
      t += 1
    when "C"
      c += 1
    when "G"
      g += 1
    end
    ans += 1 if a==t && c==g
  end
end
puts ans

n=n.to_i

s=s.chomp.split('')

ans = 0

(0...n).each do |i|

a,t,c,g = 0,0,0,0

(i...n).each do |j|

case s[j]

when "A"

a += 1

when "T"

t += 1

when "C"

c += 1

when "G"

g += 1

end

ans += 1 if a==t && c==g

end

puts ans

ジャッジ改善せず・・・

ここで、計算量は、nの2乗（今回はn=5000）のため、2.5×10の7乗です。

処理内容は1文字のStringを比較して、カウントし、最後にIntegerの比較をしているだけなので、ギリギリ間に合いそうなのに・・・

書いたコード３

こんな形で書いて、ジャッジが通りました。

カウントのループを1回で済ませる

n番目のA, T, C, Gをチェック

(2次元配列に保存)

n番目までのA, T, C, Gの個数を足し合わせる。
部分列の右と左の差を求めることで部分列の中の個数を計算する。

累積和という考え方で、毎回足すのではなく、いったん足し合わせた結果を用いて、部分列の和を求める手法です。

（計算の流れはスライド参照）

n,s=gets.split
n=n.to_i
s=s.chomp.split('')
c=Array.new(n+1){[0]*4}
(0..n).each do |i|
  case s[i]
  when ?A
    c[i+1][0]+=1
  when ?T
    c[i+1][1]+=1
  when ?C
    c[i+1][2]+=1
  when ?G
    c[i+1][3]+=1
  end
end
(1..n).each do |i|
  c[i][0]+=c[i-1][0]
  c[i][1]+=c[i-1][1]
  c[i][2]+=c[i-1][2]
  c[i][3]+=c[i-1][3]
end
ans=0
(0..n).each do |i|
  (1..n/2+1).each do |k|
    j = i+k*2
    next if i-1==j||j>n
    ai=c[j][0]-c[i][0]
    bi=c[j][1]-c[i][1]
    ci=c[j][2]-c[i][2]
    di=c[j][3]-c[i][3]
    ans+=1 if ai==bi&&ci==di
  end
end
p ans

n,s=gets.split

n=n.to_i

s=s.chomp.split('')

c=Array.new(n+1){[0]*4}

(0..n).each do |i|

case s[i]

when ?A

c[i+1][0]+=1

when ?T

c[i+1][1]+=1

when ?C

c[i+1][2]+=1

when ?G

c[i+1][3]+=1

end

(1..n).each do |i|

c[i][0]+=c[i-1][0]

c[i][1]+=c[i-1][1]

c[i][2]+=c[i-1][2]

c[i][3]+=c[i-1][3]

end

ans=0

(0..n).each do |i|

(1..n/2+1).each do |k|

j = i+k*2

next if i-1==j||j>n

ai=c[j][0]-c[i][0]

bi=c[j][1]-c[i][1]

ci=c[j][2]-c[i][2]

di=c[j][3]-c[i][3]

ans+=1 if ai==bi&&ci==di

end

p ans

想定解

https://atcoder.jp/contests/arc104/editorial/157

ほぼ、解2と同じでは・・・？

結論：やっぱり・・・Rubyは若干遅い

Goだとどうなるのか

書いたコード2と同じように書いてみました。

package main

import (
	"fmt"
	"strings"
)

func main() {
	var n int
	var s string
	sList := strings.Split(s, "")
	fmt.Scanf("%d", &n)
	fmt.Scanf("%s", &s)

	var ans int
	for i := 0; i < n; i++ {
		a, t, c, g := 0, 0, 0, 0
		for j := i; j < n; j++ {
			switch sList[j] {
			case "A":
				a++
			case "T":
				t++
			case "C":
				c++
			case "G":
				g++
			}
			if a == t && c == g {
				ans++
			}
		}
	}
	fmt.Println(ans)
}

package main

import (

"fmt"

"strings"

)

func main() {

var n int

var s string

sList := strings.Split(s, "")

fmt.Scanf("%d", &n)

fmt.Scanf("%s", &s)

var ans int

for i := 0; i < n; i++ {

a, t, c, g := 0, 0, 0, 0

for j := i; j < n; j++ {

switch sList[j] {

case "A":

a++

case "T":

t++

case "C":

c++

case "G":

g++

}

if a == t && c == g {

ans++

}

fmt.Println(ans)

}

普通にジャッジが通って、それも、Rubyで累積和を使ったコード3よりも6倍速いという結果に・・・

Rubyの一部をGoに置き換えられないか？

もっと何かしてみようということで、RubyでGoのコードを使う方法を調べて使ってみました。

RubyでGoを使うための手順

TLEになってしまったRubyのコードの一部をGoに置き換える

Go（loop.go）

// go build -buildmode=c-shared -o loop.so loop.go
package main

import (
    "C"
    "strings"
)

//export LoopCount
func LoopCount(cStr *C.char, n int) int {
    s := C.GoString(cStr)
    sSlice := strings.Split(s, "")
    var ans int
	for i := 0; i < n; i++ {
		a,t,c,g := 0,0,0,0
		for j := i; j < n; j++ {
			switch sSlice[j]{
			case "A":
				a++
			case "T":
				t++
			case "C":
				c++
			case "G":
				g++
			}
			if a == t && c == g {
				ans++
			}
		}
	}
    return ans
}

func main() {}

// go build -buildmode=c-shared -o loop.so loop.go

package main

import (

"C"

"strings"

)

//export LoopCount

func LoopCount(cStr *C.char, n int) int {

s := C.GoString(cStr)

sSlice := strings.Split(s, "")

var ans int

for i := 0; i < n; i++ {

a,t,c,g := 0,0,0,0

for j := i; j < n; j++ {

switch sSlice[j]{

case "A":

a++

case "T":

t++

case "C":

c++

case "G":

g++

}

if a == t && c == g {

ans++

}

return ans

}

func main() {}

Ruby

require 'ffi'

module Loop
  extend FFI::Library
  ffi_lib 'loop.so'
  attach_function :LoopCount, [:string, :int], :int
end


# n,s=gets.split
n = 1000
s = "ATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGA"
n = n.to_i
p Loop.LoopCount(s, n)

require 'ffi'

module Loop

extend FFI::Library

ffi_lib 'loop.so'

attach_function :LoopCount, [:string, :int], :int

end

# n,s=gets.split

n = 1000

s = "ATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGAATCGA"

n = n.to_i

p Loop.LoopCount(s, n)

スライドでは1000文字程度でしかやっていないので、若干早いかな程度でした。

n=40_000にして、時間を測ってみた

結果が以下です。（20倍くらい速い？）

# 関数部分をGoに置き換えて実行
$ time ruby ARC_104_mix.rb
300000000

real	0m5.283s
user	0m5.167s
sys	0m0.081s

# TLEになったコード
$ time ruby ARC_104_compare.rb
300000000

real	1m46.135s
user	1m40.132s
sys	0m0.406s

# 関数部分をGoに置き換えて実行

$ time ruby ARC_104_mix.rb

300000000

real 0m5.283s

user 0m5.167s

sys 0m0.081s

# TLEになったコード

$ time ruby ARC_104_compare.rb

300000000

real 1m46.135s

user 1m40.132s

sys 0m0.406s

後日談

当然Go純粋の方が圧倒的に速く、「書いたコード2」も書き方次第で通りました。

ちなみに、n=40000の時は、Go単体の場合は以下の速度が出ました。

（RubyにGoを組み込んだものと比較しても圧倒的に速い）

)$ time ./ARC_104_compare
1198

real	0m0.011s
user	0m0.007s
sys	0m0.004s

)$ time ./ARC_104_compare

1198

real 0m0.011s

user 0m0.007s

sys 0m0.004s

おしまい