HARDAtCoderLiveCodeBench

LCB-09: MEX

Baseline: PASS
Augmented: PASS
Injection: 1,981 chars

The Task

You are given a length-N sequence A=(A_1,A_2,\dots,A_N) consisting of 0, 1, and 2, and a length-N string S=S_1S_2\dots S_N consisting of M, E, and X. Find the sum of \text{mex}(A_i,A_j,A_k) over all tuples of integers (i,j,k) such that 1 \leq i < j < k \leq N and S_iS_jS_k= MEX. Here, \text{mex}(A_i,A_j,A_k) denotes the minimum non-negative integer that equals neither A_i,A_j, nor A_k. Input The input is given from Standard Input in the following format: N A_1 A_2 \dots A_N S Output Print the answer as an integer. Constraints - 3\leq N \leq 2\times 10^5 - N is an integer. - A_i \in \lbrace 0,1,2\rbrace - S is a string of length N consisting of M, E, and X. Sample Input 1 4 1 1 0 2 MEEX Sample Output 1 3 The tuples (i,j,k)\ (1 \leq i < j < k \leq N) such that S_iS_jS_k = MEX are the following two: (i,j,k)=(1,2,4),(1,3,4). Since \text{mex}(A_1,A_2,A_4)=\text{mex}(1,1,2)=0 and \text{mex}(A_1,A_3,A_4)=\text{mex}(1,0,2)=3, the answer is 0+3=3. Sample Input 2 3 0 0 0 XXX Sample Output 2 0 Sample Input 3 15 1 1 2 0 0 2 0 2 0 0 0 0 0 2 2 EXMMXXXEMEXEXMM Sample Output 3 13

Scroll to read full problem statement

The Code

Claude Opus 4.6 with max-effort extended thinking. Exact outputs. Side by side.

Baseline

PASS1257ch / 50s
import sys

def solve():
    data = sys.stdin.read().split()
    N = int(data[0])
    A = list(map(int, data[1:N+1]))
    S = data[N+1]

    # Precompute mex table
    mex_table = [[[0]*3 for _ in range(3)] for _ in range(3)]
    for a in range(3):
        for b in range(3):
            for c in range(3):
                s = {a, b, c}
                for v in range(4):
                    if v not in s:
                        mex_table[a][b][c] = v
                        break

    # Suffix counts of X positions with each value
    suffix_x = [[0]*3 for _ in range(N+1)]
    for i in range(N-1, -1, -1):
        for c in range(3):
            suffix_x[i][c] = suffix_x[i+1][c]
        if S[i] == 'X':
            suffix_x[i][A[i]] += 1

    prefix_m = [0, 0, 0]
    total = 0

    for j in range(N):
        if S[j] == 'E':
            b = A[j]
            sx = suffix_x[j+1]
            for a in range(3):
                pm = prefix_m[a]
                if pm == 0:
                    continue
                for c in range(3):
                    m = mex_table[a][b][c]
                    if m > 0 and sx[c] > 0:
                        total += pm * sx[c] * m
        if S[j] == 'M':
            prefix_m[A[j]] += 1

    print(total)

solve()

Scroll to read full code

Augmented

PASS1092ch / 148s
import sys

def solve():
    input = sys.stdin.readline
    N = int(input())
    A = list(map(int, input().split()))
    S = input().strip()

    mex_table = [[[0]*3 for _ in range(3)] for _ in range(3)]
    for a in range(3):
        for b in range(3):
            for c in range(3):
                s = {a, b, c}
                m = 0
                while m in s:
                    m += 1
                mex_table[a][b][c] = m

    countX = [0, 0, 0]
    for i in range(N):
        if S[i] == 'X':
            countX[A[i]] += 1

    countM = [0, 0, 0]
    answer = 0

    for j in range(N):
        if S[j] == 'X':
            countX[A[j]] -= 1

        if S[j] == 'E':
            e = A[j]
            for m in range(3):
                cm = countM[m]
                if cm == 0:
                    continue
                for x in range(3):
                    cx = countX[x]
                    if cx == 0:
                        continue
                    answer += cm * cx * mex_table[m][e][x]

        if S[j] == 'M':
            countM[A[j]] += 1

    print(answer)

solve()

Scroll to read full code

Blind Evaluation

Independent evaluator. Randomized A/B labels. No knowledge of which solution used the injection.

Augmented 43, Baseline 39

DimensionBaselineAugmented
Correctness1010
Efficiency78
Code Structure79
Readability78
Robustness88
Blind Analysis

Both solutions share the same core algorithm: iterate over positions where S[j] == 'E', and for each, accumulate prefix_M_count[a] * suffix_X_count[c] * mex(a, A[j], c) across all value pairs (a, c) in {0,1,2}^2. Both precompute an identical 3x3x3 mex lookup table. The O(N) outer loop with a constant 9-iteration inner loop gives O(N) time. Both are correct. Where they diverge is suffix X tracking.

Blind Observation

The baseline reasons in terms of precomputed data structures — build the suffix table, then query it. This is the instinct of someone who separates preparation from computation as a default architectural move, even when the preparation could be folded into the main pass. The augmented solution reasons in terms of loop invariants — the entire solution is organized around maintaining the correct counts at each step via a specific ordering of operations.

Source: LiveCodeBench Hard benchmark. Full code outputs: baseline · augmented · runner script & methodology

Back to Use CasesStart Building