CF 1389B - Array Walk

Rating: 1600
Tags: brute force, dp, greedy
Solve time: 1m 57s
Verified: yes

Solution

Problem Understanding

We are asked to simulate a walk along an array of positive integers, starting at the first element. Each move can either go one step to the right, which is always allowed if we are not at the last element, or one step to the left, which can only be done if we are not at the first element and if the previous move was not a left move. Each step increases a running score by the value of the array element we land on. The total number of moves must equal a given number $k$, and the total number of left moves cannot exceed $z$. The task is to maximize the score after exactly $k$ moves.

The input size can be as large as $n = 10^5$ per test case, and there can be up to $10^4$ test cases, with the total sum of $n$ across all test cases up to $3 \cdot 10^5$. This rules out any solution with worse than linear complexity per test case, because $O(n^2)$ or $O(nk)$ would result in $10^{9}$ or more operations, which is far beyond the 2-second time limit.

Edge cases that are easy to miss include arrays where $z = 0$, forcing all moves to the right, or arrays where moving left may only be useful once or twice, as excessive left moves are prohibited and two consecutive lefts are disallowed. For example, with the array [1, 100, 1], $k = 2$, and $z = 1$, the optimal sequence is right, left, yielding a score of 1 + 100 + 1 = 102. A naive greedy approach that always moves right first would only get 1 + 100 = 101, missing the optimal solution.

Approaches

The brute-force approach would generate all sequences of length $k$ that respect the left-move constraint and sum the corresponding scores. The number of sequences grows exponentially with $k$, making this completely infeasible for the given bounds.

The key insight is that the only benefit from moving left comes from pairing a left move with the immediately preceding right move, effectively forming a "left-right pair" that can be repeated at most $z$ times. Every left move must be immediately preceded by a right move, so we only need to consider how many times we can perform such pairs while moving right as much as possible. Furthermore, only the pair with the largest sum $a[i] + a[i+1]$ matters, because repeatedly using the largest adjacent pair maximizes the gain from left moves.

With this observation, the optimal strategy reduces to trying every number of left moves from 0 to $z$. For each candidate number of left moves, we calculate how many right moves are available afterward and determine the farthest index we can reach. We then compute the total score as the sum of the initial prefix plus the contributions from the repeated left-right pairs, taking care to respect the maximum allowed left moves and the total move count $k$.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(2^k)	O(k)	Too slow
Optimal	O(n * z)	O(n)	Accepted

Algorithm Walkthrough

Precompute the prefix sums of the array. This allows us to quickly calculate the total score of moving all the way to any position without left moves. Prefix sums are standard for this type of cumulative-sum problem and reduce repeated addition from $O(n)$ to $O(1)$.
Identify all adjacent pairs $a[i] + a[i+1]$ because moving left after reaching index $i+1$ allows you to repeatedly gain the sum of this pair. Track the maximum pair sum observed up to every index.
Iterate over the number of left moves $lz$ from 0 to $z$. Each left move consumes two total moves: one right move to reach a new index and one left move back. Therefore, the farthest index we can reach with $lz$ left moves is $k - 2 * lz$. For example, if $k = 7$ and $lz = 2$, the rightmost position we can initially reach is index $7 - 4 = 3$.
For each $lz$, compute the score as the sum of elements up to the farthest index reachable by right moves alone plus the sum contributed by left-right pairs. The total contribution from left moves is $lz * \text{maximum pair sum among visited positions}$.
Keep track of the maximum score across all choices of $lz$. This is the answer for that test case.

Why it works: Every left move is strictly used to benefit from the largest adjacent pair available among visited positions. Trying all counts of left moves from 0 to $z$ guarantees we consider all feasible strategies without violating constraints. The prefix sums handle all right-only moves efficiently. This exhaustively considers all optimal combinations without generating every possible sequence.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    for _ in range(t):
        n, k, z = map(int, input().split())
        a = list(map(int, input().split()))
        prefix = [0] * n
        prefix[0] = a[0]
        for i in range(1, n):
            prefix[i] = prefix[i-1] + a[i]
        max_score = 0
        max_pair = 0
        for i in range(k + 1):
            if i < n - 1:
                max_pair = max(max_pair, a[i] + a[i+1])
            left_moves = min(z, i // 2)
            right_index = i - 2 * left_moves
            if right_index < 0:
                continue
            score = prefix[right_index] + left_moves * max_pair
            max_score = max(max_score, score)
        print(max_score)

if __name__ == "__main__":
    solve()

The code first calculates prefix sums, which is used to determine the score of moving strictly to the right. The variable max_pair keeps track of the best adjacent sum for potential left moves. For each position reachable in i moves, it computes how many left moves are possible (left_moves) and calculates the total score. The score combines the prefix sum to the farthest right index and the contribution from left moves. Finally, the maximum across all feasible left-move counts is printed.

Worked Examples

Trace for input:

5 4 1
1 5 4 3 2

i (moves made)	right_index	left_moves	max_pair	prefix[right_index]	score
0	0	0	0	1	1
1	1	0	6	6	6
2	2	1	9	1 + 5 = 6	6 + 9 = 15
3	3	1	9	1 + 5 + 4 = 10	10 + 9 = 19
4	4	2	9	1 + 5 + 4 + 3 = 13	13 + 18 = 31 (not feasible, only 1 left allowed)

Max score = 19, which matches the expected output.

Trace demonstrates that the algorithm correctly limits left moves to z and calculates contributions from the optimal left-right pairs.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n * z)	Outer loop runs up to k (~n), inner computation is O(1) for left-move contribution; z ≤ 5 ensures small factor.
Space	O(n)	Prefix sums array of size n

Since z ≤ 5, effectively this is linear in n per test case. The sum of n over all test cases is ≤ 3 * 10^5, fitting comfortably within time and memory limits.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    sys.stdout = io.StringIO()
    solve()
    return sys.stdout.getvalue().strip()

# Provided samples
assert run("4\n5 4 0\n1 5 4 3 2\n5 4 1\n1 5 4 3 2\n5 4 4\n10 20 30 40 50\n10 7 3\n4 6 8 2 9 9 7 4 10 9\n") == "15\n19\n150\n56", "sample cases"

# Custom cases
assert run("1\n3 2 1\n1 100 1\n") == "102", "single left move optimal"
assert run("1\n2 1 0\n5 10\n") == "15", "minimum moves, no left allowed"
assert run("1\n5 4 2\n1 1 1 1 1\n") == "5", "all-equal values, multiple lefts allowed"
assert run("1\n5 4 5\n1 2 3 4 5\n") == "14", "z exceeds feasible left moves, only valid used