|
| 1 | +from typing import List |
| 2 | + |
| 3 | + |
1 | 4 | class Solution: |
2 | 5 | def countStableSubarrays(self, capacity: List[int]) -> int: |
3 | 6 | n = len(capacity) |
4 | | - |
5 | | - # Compute prefix sums |
| 7 | + res = 0 |
| 8 | + |
| 9 | + # Build prefix sum |
6 | 10 | prefix = [0] * (n + 1) |
7 | 11 | for i in range(n): |
8 | 12 | prefix[i + 1] = prefix[i] + capacity[i] |
9 | | - |
10 | | - res = 0 |
11 | | - |
12 | | - # For each right endpoint r, find valid left endpoints l |
13 | | - # Condition: capacity[l] == capacity[r] and |
14 | | - # capacity[l] == sum(capacity[l+1:r]) |
15 | | - # Which means: capacity[l] == prefix[r] - prefix[l+1] |
16 | | - # Rearranging: prefix[l+1] == prefix[r] - capacity[r] |
17 | | - |
18 | | - # Use a map: (value, prefix_sum) -> count |
19 | | - from collections import defaultdict |
20 | | - freq_map = defaultdict(int) |
21 | | - |
22 | | - # Process from left to right |
23 | | - for r in range(n): |
24 | | - # For each r, we want to find l such that: |
25 | | - # capacity[l] == capacity[r] and prefix[l+1] == prefix[r] - capacity[r] |
26 | | - target_prefix = prefix[r] - capacity[r] |
27 | | - key = (capacity[r], target_prefix) |
28 | | - res += freq_map[key] |
29 | | - |
30 | | - # Update map: add (capacity[r], prefix[r]) for future matches |
31 | | - # Note: we use prefix[r] (not prefix[r+1]) because when we check at position r, |
32 | | - # we're looking for l where prefix[l+1] matches |
33 | | - update_key = (capacity[r], prefix[r]) |
34 | | - freq_map[update_key] += 1 |
35 | | - |
36 | | - return res |
37 | 13 |
|
| 14 | + # For each subarray [l, r] where length >= 3 |
| 15 | + # Check if capacity[l] == capacity[r] == sum of elements between them |
| 16 | + for l in range(n): |
| 17 | + for r in range(l + 2, n): # length >= 3 |
| 18 | + if capacity[l] == capacity[r]: |
| 19 | + # Sum of elements strictly between l and r |
| 20 | + interior_sum = prefix[r] - prefix[l + 1] |
| 21 | + if capacity[l] == interior_sum: |
| 22 | + res += 1 |
| 23 | + |
| 24 | + return res |
0 commit comments