Editorial - AtCoder Weekday Contest 0171 Beta

Official

C - 旅行の計画 / Travel Plan Editorial by admin

Claude Opus 5 (High)

Overview

We are given costs for $N$ days, $A_1,\dots,A_N$. When we choose to either “cancel nothing” or “cancel a contiguous range $[l,r]$”, this problem asks us to find the total number of choices such that the remaining cost is exactly $K$ yen.

In short, we need to count the number of ranges whose contiguous subsegment sum subtracted from the total sum $S$ is exactly $T = S - K$ (plus $1$ way for “canceling nothing” if $S = K$).

Analysis

Rephrasing the Problem

Let the total sum be $S = \sum_{i=1}^N A_i$.

If we cancel nothing, the cost is $S$. Therefore, this counts as 1 way if and only if $S = K$.
If we cancel a range $[l, r]$, the cost is $S - \sum_{i=l}^{r} A_i$. This is equal to $K$ when $$\sum_{i=l}^{r} A_i = S - K$$

Thus, letting $T = S - K$, the answer we seek is:

(Number of contiguous ranges $[l,r]$ whose sum is exactly $T$) $+$ ($1$ if $S = K$, else $0$)

Note that when $T < 0$, no such range exists (since $A_i \ge 1$, any subsegment sum is strictly positive). Also, when $T = 0$, any range $[l,r]$ contains at least one day and has a sum of at least $1$, so no such range exists either (in this case, since $S = K$, the answer is just the 1 way of “canceling nothing”). Therefore, it suffices to count ranges only when $T > 0$.

Issues with the Naive Approach

Testing all pairs $(l, r)$ takes $O(N^2)$ time. For $N \le 10^6$, this requires around $10^{12}$ operations, which will not run in time. Even if we use prefix sums, enumerating the pairs themselves still takes $O(N^2)$, so the situation remains the same.

How to Solve It

The crucial observation here is the constraint that $A_i \ge 1$ (all elements are positive). Consequently, for a fixed left endpoint $l$, the range sum $f(l,r)=\sum_{i=l}^{r}A_i$ is monotonically increasing with respect to $r$; likewise, for a fixed right endpoint $r$, increasing $l$ causes the sum to monotonically decrease.

Because of this monotonicity, we can use the two pointers technique. By advancing the right pointer $r$ one step at a time and advancing the left pointer $l$ as much as necessary so that the range sum does not exceed $T$, we can quickly scan all “maximal ranges whose sum is at most $T$”.

(As an alternative approach, one could compute prefix sums $P_r$ and use a hash map to count how many times $P_r - T$ has appeared previously. While this approach also works when $A_i$ can be negative, the two pointers technique has a smaller constant factor and is much safer for $N=10^6$, which is feasible here because all elements are positive.)

Algorithm

Compute $S = \sum A_i$ and $T = S - K$.
Set $\mathrm{ans} = 0$.
If $T > 0$, count the “number of ranges whose sum is exactly $T$” using two pointers:
- Initialize $l = 0$ and $\mathrm{cur} = 0$ (the sum of the current range $[l, r]$).
- For $r = 0, 1, \dots, N-1$:
  - $\mathrm{cur} \mathrel{+}= A_r$
  - While $\mathrm{cur} > T$, repeatedly do $\mathrm{cur} \mathrel{-}= A_l$ and $l \mathrel{+}= 1$.
  - If $\mathrm{cur} = T$, then $\mathrm{ans} \mathrel{+}= 1$.
- For each $r$, the minimum $l$ such that the sum is at most $T$ is uniquely determined. If the sum there equals $T$, there is exactly one range with right endpoint $r$ whose sum is $T$ (since $A_i>0$, the range sum for a fixed right endpoint is strictly monotonic with respect to $l$, meaning at most one such $l$ exists for the same $r$). Thus, neither duplicates nor missed ranges will occur.
If $S = K$, increment $\mathrm{ans} \mathrel{+}= 1$ (the case of canceling nothing).
Output $\mathrm{ans}$.

Concrete Example

For $N = 5$, $K = 9$, $A = (1, 2, 3, 4, 5)$, we have $S = 15$ and $T = 15 - 9 = 6$.

Looking for ranges with sum $6$, we find $[1,3]$ ($1+2+3$) and $[4,4]$ ($A_4 = 6$? No, $A_4=4$, so not this one)… in reality, only $[1,3]$ from $\{1,2,3\}$ matches. Trace of the two pointers method:

$r$	cur (after adding $A_r$)	$[l,r]$ and cur after shrinking	Check
1	1	$[1,1]$, 1	–
2	3	$[1,2]$, 3	–
3	6	$[1,3]$, 6	Match
4	10 → 9 → 7 → 4	$[4,4]$, 4	–
5	9 → 5	$[5,5]$, 5	–

Therefore, there is 1 such range. Since $S \neq K$, the answer is $1$.

Complexity

Time complexity: $O(N)$ (since both $l$ and $r$ advance monotonically at most $N$ times each)
Space complexity: $O(N)$ (to store the input array; can be $O(1)$ if elements are processed on the fly without storing them)

Implementation Notes

Handling $T < 0$ and $T = 0$: When $T \le 0$, no valid range exists, so we skip the two pointers scan. In particular, beware of the bug where an “empty range” might be counted when $T = 0$ (the problem specifies ranges with $l \le r$, meaning at least one day must be canceled).
Don’t forget the “cancel nothing” case: When $S = K$, add $+1$ independently of counting ranges. While this corresponds exactly to $T = 0$, it is safer to treat it separately from the range count.
Integer Overflow: $S$ can be up to $10^6 \times 10^9 = 10^{15}$ and $K$ can also be up to $10^{15}$, so 64-bit integers are required (long long in C++). In Python, integers have arbitrary precision, so overflow is not a concern.
Fast I/O: Since $N \le 10^6$, in Python it is important to read the input in bulk using sys.stdin.buffer.read().split(). Calling input() $10^6$ times will cause TLE. In C++, using scanf or fast I/O is also recommended.
Because we do not explicitly build a prefix sum array and only add/subtract from cur, this approach is advantageous in terms of both memory and constant factors.

Source Code

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    k = int(data[1])
    A = [0] * n
    for i in range(n):
        A[i] = int(data[2 + i])
    S = sum(A)
    T = S - k
    ans = 0
    if T > 0:
        l = 0
        cur = 0
        for r in range(n):
            cur += A[r]
            while cur > T:
                cur -= A[l]
                l += 1
            if cur == T:
                ans += 1
    if S == k:
        ans += 1
    sys.stdout.write(str(ans) + "\n")

main()

This editorial was generated by claude5opus-high.

\(r\)	cur (after adding \(A_r\))	\([l,r]\) and cur after shrinking	Check
1	1	\([1,1]\), 1	–
2	3	\([1,2]\), 3	–
3	6	\([1,3]\), 6	Match
4	10 → 9 → 7 → 4	\([4,4]\), 4	–
5	9 → 5	\([5,5]\), 5	–