Posted by まえおき
最初に記事の本題に入る前に前置きを書きたいと思います。去年くらいから新しくプログラミング言語を勉強していても実務で使用する機会がない言語でアウトプットする場がないと身につかないなと思い始めました。そこで色々考えていたのですが、ある時TwitterのTLでAtCoder(https://atcoder.jp/)という日本発の競技プログラミングが流行りはじめているのをみて興味を持ったのがきっかけでした。試しに参加してみたら割と簡単に参加できて、大体週一回開催というペースがあっているなと感じ少しづつ始めるようになりました。言語は最近勉強し始めていたPythonを使って参加しておりますが、数回参加して感じたのはプログラミング言語の問題や課題ではなく自分のアルゴリズム能力の無さでした。この課題は将来のためにも改善した方が良いと思い、より本腰を入れて毎週末の夜にABCのコンテストに参加するようになりました。ところが、参加するだけではなかなか能力の向上はないので、問題に臨んだけど解くことができなかった問題は回答や他の方の実装を参照して、自分なりの実装を作成するというのを始めることにしました。そして、その内容はなるべくブログで簡単にでもアウトプットしようと思っています。それの初回が本記事となります。また、時間があれば、アルゴリズムについてチーター本、蟻本、螺旋本あたりを読んでブログで整理できればと考えています。ということで、前置きが長くなりましたが、早速先週末(4/12)の夜に開催されたABC162のD問題が解けなかったので、その回答と自分なりの実装を本記事では記載します。
問題
ABC162 D問題 RGB Triplets
https://atcoder.jp/contests/abc162/tasks/abc162_d
解説
https://img.atcoder.jp/abc162/editorial.pdf
理解
何も考えないで実装すると、i, j, kについて全探索で処理を行ってしまい計算量はO(N^3)となり、nは最大で4000なのでTLEしてしまいそうです。なので、方針として最初に”R”,”G”,”B”の3つが異なる組み合わせとなる総数を算出してから2つの条件にマッチしないものを引いていくことにします。その理由は、マッチしないパターンであれば問題文の2個目の条件でi, j, kの関係でi, jを固定するとkが算出できるので、その分を総数から引いてあげることができます。その場合の計算量はO(N^2)とすることができそうです。
実装例
自分が作成した実装は下記の通りとなります。
N = int(input())
S = list(input())
R = 0
G = 0
B = 0
for i in range(N):
if S[i] == 'R':
R += 1
elif S[i] == 'G':
G += 1
else:
B += 1
ans = R * G * B
for i in range(N - 2):
for j in range(i + 1, N - 1):
k = j + (j - i)
if k >= N:
break
if S[i] != S[j] and S[i] != S[k] and S[j] != S[k]:
ans -= 1
print(ans)