競プロ初中級者100問

JOI2008本戦A 碁石ならべ

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/joi2008ho/tasks/joi2008ho_a
https://www.ioi-jp.org/joi/2007/2008-ho-prob_and_sol/2008-ho.pdf#page=2

方針

愚直にルールにしたがってシミュレーションを行うと、計算量はO(n^2)となり間に合いません。

少し考察すると色がどこで変わったかだけを記録すると効率が良いとわかります。

これはスタックで実装することができ、計算量はO(n)です。

コード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
N = int(input())
C = []
for n in range(N):
    c = int(input())
    C.append(c)

stack = []
old = None
for n in range(N):
    c = C[n]
    if (n%2==1): # iが偶数の場合（nが奇数）
        if old != c: # 色が異なれば
            stack.pop() # 反転させる（最後の反対色の記録を取り除く）
            if len(stack)==0: # スタックを空にしない
                stack.append(0)
    else:
        if old != c: # 違う色が出た時
            stack.append(n) # 色の始まりを記録する
    old = c
if old == 0: # スタック上の記録では、白黒は交互に格納されているので最後が白なら
    stack.append(n+1) # n+1を格納しておく。これは、後ほどストライド2で差分を計算するため
if len(stack)%2 == 1:
    stack.insert(0,0) # ストライド2で差分を計算するため、要素数を偶数にする
print (sum(stack[1::2]) - sum(stack[::2]))

記事情報

  • 投稿日：2020年5月8日
  • 最終更新日：2020年5月8日