Ruby と Python と Java で解く AtCoder ARC104 B 累積和

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder Regular Contest B - DNA Sequence
Difficulty: 410

今回のテーマ、累積和

累積和問題ですが、今回はACGTの4個の累積和が必要になります

Ruby

ruby.rb

n, s = gets.chomp.split
n = n.to_i
a = [0]
c = [0]
g = [0]
t = [0]
n.times do |i|
  a << a[-1]
  c << c[-1]
  g << g[-1]
  t << t[-1]
  if s[i] == 'A'
    a[i + 1] += 1
  elsif s[i] == 'C'
    c[i + 1] += 1
  elsif s[i] == 'G'
    g[i + 1] += 1
  elsif s[i] == 'T'
    t[i + 1] += 1
  end
end
cnt = 0
n.times do |i|
  at = a[i] - t[i]
  cg = c[i] - g[i]
  (i.next).upto(n) do |j|
    cnt += 1 if a[j] - t[j] == at && c[j] - g[j] == cg
  end
end
puts cnt

コメント欄の @scivola さんのアイデアを反映させていただきました。

ruiseki.rb

  a << a[-1]
  c << c[-1]
  g << g[-1]
  t << t[-1]

累積した配列を作る場合、最初に必要な要素数を用意することもできます。
今回は[-1]で最終要素を呼び出し追加してみました。

hantei.rb

  at = a[i] - t[i]
  cg = c[i] - g[i]
  (i.next).upto(n) do |j|
    cnt += 1 if a[j] - t[j] == at && c[j] - g[j] == cg
  end

区間内のACGTの要素数が同じであれば相補的と判断できます。

Python

pypy.py

from sys import stdin

def main():
    input = stdin.readline
    n, s = input().split()
    n = int(n)
    s = 'N' + s + 'N'
    cnt = 0
    for i in range(1, n + 2):
        a = 0
        c = 0
        g = 0
        t = 0
        j = i + 1
        while True:
            if s[i] == 'A':
                a += 1
            elif s[i] == 'C':
                c += 1
            elif s[i] == 'G':
                g += 1
            elif s[i] == 'T':
                t += 1
            else:
                break
            if s[j] == 'A':
                a += 1
            elif s[j] == 'C':
                c += 1
            elif s[j] == 'G':
                g += 1
            elif s[j] == 'T':
                t += 1
            else:
                break
            if a == t and c == g:
                cnt += 1
            i -= 1
            j += 1
    print(cnt)
main()

pypyは別の解法です。
この解法ですとPythonやRubyではTLEになります。

banpei.py

    s = 'N' + s + 'N'

文字列の前後に番兵を置いています。
whileでbreakする条件に使用します。

hiroge.py

            i -= 1
            j += 1

各文字位置で範囲を広げて個数を確認しています

Java

java.java

import java.util.*;

class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = Integer.parseInt(sc.next());
        String s = "N" + sc.next() + "N";
        sc.close();

        int cnt = 0;
        for (int k = 1; k < n+2; k++) {
            int i = k;
            int j = k + 1;
            int gn = 0;
            int cn = 0;
            int tn = 0;
            int an = 0;
            while (true) {
                if (s.charAt(i) == 'A') {
                    an++;
                } else if (s.charAt(i) == 'C') {
                    cn++;
                } else if (s.charAt(i) == 'G') {
                    gn++;
                } else if (s.charAt(i) == 'T') {
                    tn++;
                } else {
                    break;
                }
                if (s.charAt(j) == 'A') {
                    an++;
                } else if (s.charAt(j) == 'C') {
                    cn++;
                } else if (s.charAt(j) == 'G') {
                    gn++;
                } else if (s.charAt(j) == 'T') {
                    tn++;
                } else {
                    break;
                }
                if (an == tn && cn == gn) {
                    cnt++;
                }
                i--;
                j++;
            }
        }
        System.out.println(cnt);
    }
}

JavaはPypyと同じ解法です。
コメント欄の @c-yan さんのアイデアを反映させていただきました。

i.java

        for (int k = 1; k < n+2; k++) {
            int i = k;

PythonやRubyとは異なり、ブロック内でループ変数に代入を行うと、次回のループ変数に影響が出ますので、別の変数に代入しています。

	Ruby	Pypy	Java
コード長 (Byte)	482	920	1436
実行時間 (ms)	1011	336	335
メモリ (KB)	14676	69312	39456

まとめ

• ARC 104 D を解いた
• Ruby に詳しくなった
• Python に詳しくなった
• Java に詳しくなった