Official

B - サンドイッチメロディ / Sandwich Melody Editorial by admin

Gemini 3.0 Flash (Thinking)

概要

この問題は、与えられた文字列 \(S\) の連続部分文字列のうち、その「圧縮列」が長さ \(3\) かつ \(1\) 文字目と \(3\) 文字目が一致するものの個数を求める問題です。

考察

圧縮列の性質と「サンドイッチ区間」の条件

文字列の圧縮列とは、連続する同じ文字を \(1\) 文字にまとめたものです。 例えば、AAABBAAAA の圧縮列は ABA となります。

問題文にある「サンドイッチ区間」の条件を整理すると、連続部分文字列 \(S[l..r]\) の圧縮列が \(X_1 X_2 X_3\) という形になり、かつ \(X_1 = X_3\) である必要があります。 これは、部分文字列が以下のような構造をしていることを意味します。

  1. 文字 \(A\)\(1\) 回以上続く(ブロック1)
  2. 文字 \(B\) (\(B \neq A\)) が \(1\) 回以上続く(ブロック2)
  3. 文字 \(A\)\(1\) 回以上続く(ブロック3)

そして、この部分文字列の直前や直後に他の文字が含まれていても、圧縮列の長さが「ちょうど \(3\)」であるためには、選んだ範囲 \([l, r]\)ちょうど \(3\) つの異なる文字のブロックにまたがっている必要があります。

効率的な数え上げ

\(N\) が最大 \(10^6\) であるため、すべての区間 \((l, r)\) を全探索すると \(O(N^2)\) かかり、間に合いません。そこで、文字列を同じ文字が連続する「ブロック」に分割して考えます(ランレングス圧縮)。

例えば、AAABBBAA は以下の \(3\) つのブロックに分けられます。 - ブロック 0: A\(3\) 個 - ブロック 1: B\(3\) 個 - ブロック 2: A\(2\)

ここで、圧縮列が ABA となるような区間 \([l, r]\) は、「ブロック 0 のいずれかの位置から始まり、ブロック 1 をすべて含み、ブロック 2 のいずれかの位置で終わる」ような区間に他なりません。

一般に、連続する \(3\) つのブロック \(i, i+1, i+2\) において、ブロック \(i\) とブロック \(i+2\) の文字が同じである場合、以下の条件を満たす \((l, r)\) はすべてサンドイッチ区間になります。 - \(l\) はブロック \(i\) 内のどの位置でもよい。 - \(r\) はブロック \(i+2\) 内のどの位置でもよい。

このとき、ブロック \(i\) の長さを \(L_i\)、ブロック \(i+2\) の長さを \(L_{i+2}\) とすると、その組み合わせは \(L_i \times L_{i+2}\) 通り存在します。

アルゴリズム

  1. ランレングス圧縮: 文字列 \(S\) を走査し、同じ文字が連続する区間を「文字」と「その長さ」のリストに変換します。 例: AAABBBAA \(\rightarrow\) [('A', 3), ('B', 3), ('A', 2)]
  2. 三つ組の判定: 圧縮されたブロックのリストを先頭から順に見ていきます。 連続する \(3\) つのブロック \((i, i+1, i+2)\) について、ブロック[i].文字 == ブロック[i+2].文字 であるかを確認します。
  3. 個数の加算: 条件を満たす場合、ブロック[i].長さ * ブロック[i+2].長さ を答えに加算します。
  4. 結果の出力: 最終的な合計値を出力します。

計算量

  • 時間計算量: \(O(N)\) 文字列を \(1\) 回走査してブロック化するのに \(O(N)\)、ブロックのリスト(最大 \(N\) 個)を \(1\) 回走査するのに \(O(N)\) かかるため、全体で \(O(N)\) です。
  • 空間計算量: \(O(N)\) ブロック化したデータを格納するためのメモリが必要です。

実装のポイント

  • 文字列の長さ \(N\) が大きいため、Python では sys.stdin.read().split() などを用いて高速に入力を読み込むのが望ましいです。

  • 隣接するブロックは必ず異なる文字になるため、ブロック[i].文字 == ブロック[i+2].文字 が成り立つとき、自動的に真ん中の ブロック[i+1] は異なる文字であることが保証されます。

    ソースコード

import sys

def solve():
    # Read all input from standard input at once
    input_data = sys.stdin.read().split()
    if len(input_data) < 2:
        return
    
    # The first element is N, the length of the string
    n = int(input_data[0])
    
    # The string S may be split across multiple lines in some input formats,
    # so we join all subsequent elements to form the complete string.
    # For N=10^6, this is efficient enough in Python.
    s = "".join(input_data[1:])
    
    # Update n to be the actual length of the string s
    n = len(s)
    
    # A sandwich interval must have a compressed sequence of length exactly 3.
    # This requires the substring to span at least 3 characters.
    if n < 3:
        print(0)
        return

    # To find all sandwich intervals, we first compress the string into blocks
    # of identical consecutive characters. Each block is represented by its
    # character and its length.
    chars = []
    lengths = []
    
    if n > 0:
        curr_char = s[0]
        curr_len = 0
        for char in s:
            if char == curr_char:
                curr_len += 1
            else:
                chars.append(curr_char)
                lengths.append(curr_len)
                curr_char = char
                curr_len = 1
        # Add the final block
        chars.append(curr_char)
        lengths.append(curr_len)
    
    num_blocks = len(chars)
    ans = 0
    
    # A sandwich interval corresponds to a sequence of exactly three consecutive
    # blocks (B_i, B_{i+1}, B_{i+2}) such that the character in the first block
    # is the same as the character in the third block (chars[i] == chars[i+2]).
    #
    # Because adjacent blocks always have different characters, if chars[i] == chars[i+2],
    # the character in the middle block (chars[i+1]) must be different.
    # Thus, the compressed sequence for any substring starting in block i,
    # covering all of block i+1, and ending in block i+2 will be
    # (chars[i], chars[i+1], chars[i]). This sequence has length exactly 3
    # and its 1st and 3rd characters are equal.
    #
    # The number of such intervals for a specific triple of blocks is the
    # number of ways to choose a starting position in the first block (lengths[i])
    # multiplied by the number of ways to choose an ending position in the
    # third block (lengths[i+2]).
    for i in range(num_blocks - 2):
        if chars[i] == chars[i+2]:
            ans += lengths[i] * lengths[i+2]
            
    # Output the total count of sandwich intervals
    print(ans)

if __name__ == '__main__':
    solve()

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

posted:
last update: