このシリーズではE869120さんによって作成された競プロ典型90問をPythonで解いています。
問題
084 - There are two types of characters（★3）
ポイント
o と x 両方が含まれる選び方を考えていくと時間がかかりそうなので、余事象の考え方を使うことにします。
すなわち
（全ての選び方）- （oしか含まれない選び方）-（×しか含まれない選び方）を求めていきます。
先に全ての選び方について考えます。
lとrは重なっても良い（1≤l≤r≤N）という条件なので、rが1文字目の場合はlも1文字目しか選べません。rが2文字目であれば、lの選び方は1文字目と2文字目の2通りです。rが3文字目であれば同様に3通り…となります。
全ての選び方はrが1文字目~N文字目の場合のlの選び方の合計に等しいので、等差数列の公式を
使って、N * (1+N) / 2 で求めることができます。
次に、oしか含まれない、xしか含まれない選び方を考えましょう。
前準備として、oしか含まれない、xしか含まれない区間の情報を集めます。
そのために、文字列Sの先頭から同じ文字が続いている区間ごとにその文字数を記録していきます。
例えば文字列Sが xxooxoxooo であった場合、同じ文字が続いている区間の文字数は
2,2,1,1,1,3となります。
※このように同じ文字が連続する区間についてその文字数を記録していく方法をランレングス圧縮といいます。
この時、各区間ごとのl,rの選び方の総数 = oしか含まれない、xしか含まれない選び方の総数 です。選び方は同じように等差数列で求めることができます。
以上を踏まえ、実装したコードがこちらです。
コード
def main():
import sys
sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline
N = int(input())
S = input()[:-1]
block = []
num = 0
# ランレングス圧縮
for i in range(N):
num += 1
if i != N - 1 and S[i] != S[i + 1]:
block.append(num)
num = 0
block.append(num) # 上の処理だと最後の区間をblockに入れられないので、ここで最後の区間をblockに入れる。
all = N * (1 + N) // 2 # 全ての選び方
# oしか含まれない、xしか含まれない選び方
com = 0
for b in block:
com += b * (1 + b) // 2
print(all - com)
if __name__ == '__main__':
main()
