ARC104 B - DNA SequenceをSwiftで解く！

はじめに

こんにちは！屋代高校PC同好会結成会のTARDIGRADEです。
本記事ではARC104 B - DNA SequenceをSwiftで解いていきたいと思います！
2020/10/24に開かれたSwift Zoomin' #4に参加された方向けの解説になります。
競技プログラミングを始めて体験した方にもわかりやすいように、問題を読んだ時からコーディングするまでの思考の過程も伝えられたらと思います。
深夜に急いで書いたので解説が雑な部分もありますが、あらかじめご了承ください。質問等あれば、コメントかTwitterの方にお願いします。

Step1：問題を理解する

当然ですが、問題を解くには問題の内容をしっかり理解しなければなりません。競技プログラミングの問題は数学のテストのような形で書かれることも多く、初めて見る方は「うっ…」となるかもしれませんが、頑張って読み進めていきましょう。

問題を理解するときのコツ

・自分で図を書いたり式を書いたりして問われていることを可視化してみる
可視化することで効率のいい解法がぱっと思いつくことも多いです。

・テストケースをよく見る
テストケースにプログラムの動く流れが記載されていることがあります。問題の理解を進めるうえで非常に参考になるので、しっかりと読みましょう。また、複数のテストケースを見ることで「複雑に見えるけど、実際答えはAかBの二択しかない！」というように考察が進められることもあります。

B - DNA Sequenceで問われていること

「相補的」という言葉が結構なパワーワードで拒絶反応を起こしやすいですが、よく読むと問題文に「相補的」の意味が書かれています。

$|T1|=l$としたとき、どの $1≤i≤l$についても $T1,T2$の$i$文字目の組み合わせが $(A$と$T)$, または$(C$と$G)$ の組み合わせのいずれかであることを指します。

|T1|=lとしたとき、どの 1≤i≤lについてもの部分は定型文みたいなものなので、無視していただいて結構です。定義をきちんとしないと文句を言われてしまうので、問題文はめちゃめちゃ細かく丁寧に書かれますが、「大したこと言ってないじゃん」という部分も多いです。この辺は数学とよく似ていますね。

さて、T1,T2のi文字目の組み合わせが (AとT), または(CとG) の組み合わせのいずれかであることを指しますという部分だけでしたら理解はそんなに難しくないでしょう。問題文を読み進めると、ある部分文字列$T$を並べ替えたときに、$T$と相補的な文字列が存在するかどうかを判定しなくてはいけないことがわかります。相補的な文字列であるとはつまり、もともと$"A"$であったところは$"T"$、$"T"$であったところは$"A"$に代わっているということです。$(C$と$G)$についても同様のことが言えます。

ここまで考察を進めると、判定の方法についてとあることに気づきます。というか気づけるかどうかがこの問題の鍵になってくるので、あまりピンと来ていない方はしばらく考えてみましょう。

………そうです！部分文字列$T$を並べ替えたときに、$T$と相補的な文字列が存在するためには、$T$に含まれている $"A"$と$"T"$ の数が等しく、かつ $"C"$と$"G"$ の数が等しくなければなりません。

これでこの問題は、「$S$ の連続する空でない部分文字列 $T$ であって、$T$に含まれている$"A"$と$"T"$ の数が等しく、かつ $"C"$と$"G"$ の数が等しい物の個数を求めよ」という問題に置き換えることができます。

このように、競プロでは「問題文をわかりやすい形に置き換える」ことも重要になります。

Step2：愚直な解法を考える

問題を簡潔にまとめられたので、次に解法を考えていきます。この時非常に参考になるのが標準入力で与えられる値の制約です。
今回の場合$N <= 5000$ですので、$O(N^2)$解法が間に合うことになります (計算量オーダーの考え方がわからない方はこちらのサイトを見てください) 。

さて、計算量の上限もわかったことですし、これから本格的に解法を考えていきましょう。$T$の中にある$"A","T","C","G"$ の数がわかれば、$T$に相補的な文字列があるかどうか判定できることがStep1で分かっているので、とりあえずはそれを愚直に書いてみます。

let input = readLine()!.split(separator: " ")
let n = Int(input[0])!
let s = Array(input[1]).map{String($0)}
var answer = 0
for i in 0..<n-1{
    for j in i+1..<n{
        var ac = 0
        var tc = 0
        var cc = 0
        var gc = 0
        for k in i...j{
            if s[k] == "A"{ac += 1}
            if s[k] == "T"{tc += 1}
            if s[k] == "C"{cc += 1}
            if s[k] == "G"{gc += 1} 
        }
        if ac == tc && cc == gc{answer += 1}
    }
}
print(answer)

実際にテストケースを入れてみると、これで正しい答えを求められていることがわかります。
ただし、このプログラムの計算量は$O(N^3)$であり、これだと実行時間制限に引っかかってしまいます (実際に提出してみるとこんな感じ) 。このように、先に計算量の上限を見積もっておけば、TLEになるコードを提出してペナルティを食らうことを防げます。

Step3：プログラムの高速化

ここから、どのように高速化すればいいのかを考えてみます。
そもそも現状では、部分文字列の指定で$O(N^2)$かかってしまっているので、「相補的な文字列が存在するかどうかの判定を$O(1)$でする」か、もしくは「部分文字列を個別に指定せずに解く方法」を考えなければなりません。今回の問題では前者について考えるとうまくいきそうです。「どの点に着目して高速化するべきか」、というのは量をこなせばだんだんとわかってくるので、今はわからなくても心配しないでください。

前者の解法はこちら。クリックして展開できます。

さて、判定を$O(1)$でやらなければならないわけですが、いったいどうすればいいでしょう。
ここで、「判定を$O(1)$で行う」ということについてもう少し詳しく考えてみましょう。ここでいう判定には、$"A","T","C","G"$ の数が必要です。具体的には$S[i]$から$S[j]$に含まれている各文字の数です。
今は$i$から$j$までループを回して数えていますが、ここを$O(1)$で取得できないでしょうか…？もしできれば判定自体も$O(1)$でできることになり、全体の計算量は$O(N^2)$に改善されます。

発想を飛ばしやすくするために、$S[i]$から$S[j]$に含まれているそれぞれの文字の数を求める式を変形してみましょう。

(求めたいもの) = ($S[i]$から$S[j]$に含まれているそれぞれの文字の数)

= $(S[i]$に含まれている各文字の数$)$ + $(S[i+1]$に含まれている各文字の数$)$ + ...... + $(S[j-1]$に含まれている各文字の数$)$ + $(S[j]$に含まれている各文字の数$)$

= $(S[1]$から$S[j]$に含まれている各文字の数$)$ - $(S[1]$から$S[i-1]$に含まれている各文字の数$)$

ここまで変形すると、だいぶ式がシンプルになりました。そして最後の式をよく見てください。これは累積和というものを使うことで$O(1)$で求めることができます。
累積和については、僕が解説するよりもほかにずっとわかりやすい記事があるのでそちらを紹介しておきます。

累積和を何も考えずに書けるようにする！

これを使えば、判定を$O(1)$で、全体の処理を$O(N^2)$で行うことができます！
実装すると以下のようになります。
(この辺の解説が雑で申し訳ないのですが、一番重要なところなので、リンク先の記事等を参考に、やっていることを理解できるようにしましょう)

let input = readLine()!.split(separator: " ")
let n = Int(input[0])!
let s = input[1]
var aa:[Int] = [0]
var tt:[Int] = [0]
var cc:[Int] = [0]
var gg:[Int] = [0]
var a = 0
var t = 0
var c = 0
var g = 0
// 累積和を使うための前処理
for i in s{
    if i == "A"{a += 1}
    if i == "T"{t += 1}
    if i == "C"{c += 1}
    if i == "G"{g += 1}
    aa.append(a)
    tt.append(t)
    cc.append(c)
    gg.append(g)
}
// ここからは先ほど作成した愚直な方法と同じ
var ans = 0
for i in 0...n-1{
    for j in i+1...n{
        if aa[j]-aa[i] == tt[j]-tt[i] && cc[j]-cc[i] == gg[j]-gg[i]{
            ans += 1
        }
    }
}
print(ans)

2020/10/26　追記
問題を読み直したところ、どちらでもできることがわかりました。というかむしろ、考え方的にも実装量的にも後者の方が簡単です。

具体的に、後者の場合はどのような考え方で解いていくのかを説明したいと思います。
まず、Step2で書いたプログラムがどのように動いていくのかを眺めてみましょう。すると、計算に無駄があることに気づくと思います。具体的には↓の部分です。

for j in i+1..<n{
    //略   
    for k in i...j{
        //略
    }

一回ずつ丁寧にプログラムが動く流れを追ってみると、

$i...j$の値をループで求める
$→i...j+1$の値をループで求める
$→i...j+2$の値をループで求める
$→$以下同様に繰り返す

というように動いていることがわかります。
しかし、よく考えてみると、 $i...j+1$ の値というのは $i...j$ の値に $j+1$ の値を足すだけで求められます。わざわざ $i...j$ の値をもう一度求め直すというのは無駄な計算です。$i...j+2$ の値を求める時も同様に、 $i...j+1$ の値に $j+2$ の値を足すだけでよいのです。

つまり何が言いたいのかというと、一度計算した値は、次の値を求めるのに利用できるということです。ここを改善すると、累積和なんて使わなくても、計算量を$O(N^2)$まで落とせます。
実装してみると以下のようになります。

let input = readLine()!.split(separator: " ")
let n = Int(input[0])!
let s = Array(input[1]).map{String($0)}
var answer = 0
for i in 0..<n-1{
    var ac = 0
    var tc = 0
    var cc = 0
    var gc = 0
    for j in i..<n{
        if s[j] == "A"{ac += 1}
        if s[j] == "T"{tc += 1}
        if s[j] == "C"{cc += 1}
        if s[j] == "G"{gc += 1} 
        if ac == tc && cc == gc{answer += 1}
    }
}
print(answer)

こうやって無駄な計算をしてしまっている部分を抜き出してみると「そりゃそうじゃん」という話になるんですが、慣れてないとなかなか気づきません。はじめのうちは、「愚直な解法を考える」→「プログラムが動く流れを見て無駄な計算をしている部分がないか探す」という感じで解くことをお勧めします。

最後に

解説は以上で終わりとなります。説明がわかりにくいところも多かったかもしれませんが(僕の力不足です、すみません)、問題を解くときの大体の雰囲気が伝わってくれていればいいなと思っております。
冒頭にも書きましたが、何か質問等ございましたら、コメントかTwitterの方で受け付けております。