Official

A - 投票の行方 / The Fate of the Vote Editorial by admin

gemini-3.1-pro-thinking

概要

生徒の支持状態が反転するイベントを順に処理し、初めてどちらかの候補者の支持者が \(0\) 人になるタイミングを求める問題です。

考察

演説が行われるたびに生徒の支持状態が変化しますが、毎回すべての生徒の支持状態を確認して高橋君と青木君の支持者数を数え直すと、1回の確認に \(O(N)\) の時間がかかります。これを \(M\) 回繰り返すと全体の計算量が \(O(NM)\) となり、制約(\(N, M \leq 2 \times 10^5\))のもとでは実行時間制限(TLE)を超えてしまいます。

これを解決するためには、「全体の支持者数」をあらかじめ変数として持っておき、変化があった差分だけを更新するというアプローチをとります。

例えば、ある生徒の支持が「高橋君(T)」から「青木君(A)」に変わったとします。このとき、全体の支持者数は以下のように変化します。 - 高橋君の支持者数:\(1\) 人減る - 青木君の支持者数:\(1\) 人増える

このように、演説のたびに2つの変数を \(+1\) または \(-1\) するだけであれば、1回あたりの処理が \(O(1)\) で完了します。あとは、各生徒が現在どちらを支持しているかを記録した配列を更新しながらシミュレーションを行えば、高速に答えを求めることができます。

アルゴリズム

  1. 初期状態の文字列 \(S\) を、文字の変更が可能な配列(リスト)に変換します。
  2. 初期状態における高橋君の支持者数 t_count と、青木君の支持者数 a_count を数えます。
  3. \(M\) 回の演説について、順番に以下の処理を行います。
    • 対象となる生徒 \(R_i\) の現在の支持状態を確認します。
    • T の場合:配列の値を A に書き換え、t_count\(1\) 減らし、a_count\(1\) 増やします。
    • A の場合:配列の値を T に書き換え、a_count\(1\) 減らし、t_count\(1\) 増やします。
    • 変更の直後に t_count または a_count\(0\) になっているかを確認します。もし \(0\) なら、その時点での演説の回数(\(i\) 回目)を出力してプログラムを終了します。
  4. \(M\) 回の処理を終えてもどちらの支持者数も \(0\) にならなかった場合は、\(-1\) を出力します。

計算量

  • 時間計算量: \(O(N + M)\)
    • 初期状態のカウントや配列の作成に \(O(N)\) かかります。
    • 各演説の処理は \(O(1)\) で行えるため、\(M\) 回の処理で \(O(M)\) かかります。
    • したがって、全体で \(O(N + M)\) となり、十分に高速です。
  • 空間計算量: \(O(N + M)\)
    • 生徒の支持状態を管理する配列に \(O(N)\) のメモリを使用します。
    • 入力データを一括で読み込むためのリストに \(O(N + M)\) のメモリを使用します。

実装のポイント

  • 文字列の変更可能性(イミュータブル) Pythonの文字列(str 型)は一部の文字だけを直接書き換えることができません(例: S[0] = 'A' はエラーになります)。そのため、最初に list(S) を使って文字のリストに変換しておく必要があります。

  • インデックスの調整(0-indexed と 1-indexed) 入力で与えられる生徒の番号 \(R_i\)\(1\) から始まる「1-indexed」ですが、Pythonの配列のインデックスは \(0\) から始まる「0-indexed」です。配列にアクセスする際は、入力値から \(1\) を引いて idx = R_i - 1 として扱うことに注意しましょう。

  • 高速な入出力 入力データが多いため、sys.stdin.read().split() を用いて一括で読み込むことで、入出力にかかる時間を短縮しています。

    ソースコード

import sys

def main():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    M = int(input_data[1])
    S = input_data[2]
    
    state = list(S)
    t_count = S.count('T')
    a_count = N - t_count
    
    for i in range(M):
        idx = int(input_data[3 + i]) - 1
        if state[idx] == 'T':
            state[idx] = 'A'
            t_count -= 1
            a_count += 1
        else:
            state[idx] = 'T'
            a_count -= 1
            t_count += 1
            
        if t_count == 0 or a_count == 0:
            print(i + 1)
            return
            
    print(-1)

if __name__ == '__main__':
    main()

この解説は gemini-3.1-pro-thinking によって生成されました。

posted:
last update: