1. 문제 소개

Constraints:

• 1 <= k <= nums.length <= 2 * 105
• 109 <= nums[i] <= 109

3381. Maximum Subarray Sum With Length Divisible by K

2. 나의 풀이

혼자 고민해봤을 때 prefix를 사용하는 방법까지는 떠올랐지만, mod를 잘 사용하지 못했다. 힌트를 봐가며 풀었다.
시간복잡도는 $O(n)$이다.

class Solution:
    def maxSubarraySum(self, nums: List[int], k: int) -> int:
        n = len(nums)
        INF = 10 ** 15                      # 충분히 큰 양수 (최소값 초기화용)
        pref = [0] * (n + 1)                # prefix[i] = nums[0..i-1] 의 합

        for i in range(1, n+1):
            pref[i] = pref[i-1] + nums[i-1] # prefix 합 갱신: 이전 prefix + 현재 원소

        # min_pref[mod] = 지금까지 본 prefix 인덱스 중,
        #   인덱스 % k == mod 인 것들 중 가장 작은 prefix 값
        min_pref = [INF] * k
        min_pref[0] = 0                     # prefix[0] = 0, 인덱스 0 % k == 0

        ans = -INF                          # 가능한 합들보다 작은 값으로 초기화

        for idx in range(1, n+1):
            mod = idx % k                   # 현재 prefix 인덱스의 mod k 그룹

            # min_pref[mod] 가 갱신된 적이 있다면 (즉, 유효한 prefix가 있다면)
            # prefix[idx] 를 오른쪽 끝으로 하는, 길이가 k의 배수인 subarray 후보를 구할 수 있다.
            # 해당 subarray 합 = prefix[idx] - min_pref[mod]
            if min_pref[mod] != INF:
                ans = max(ans, pref[idx] - min_pref[mod])

            # 현재 prefix[idx] 를 이용해 이 mod 그룹의 최소 prefix 값 갱신
            min_pref[mod] = min(min_pref[mod], pref[idx])
    
        return ans                          # 길이가 k의 배수인 subarray 중 최대 합 반환

3. 다른 풀이

가장 빠른 풀이를 가져왔다.
역시 max(), min() 보다는 부등호가 훨씬 빠르다.

from math import inf

class Solution:
    def maxSubarraySum(self, nums: List[int], k: int) -> int:
        least_pre = [inf] * k                # 각 mod 그룹별로 '가장 작은 prefix sum'을 저장하는 배열
        least_pre[-1] = 0                    # prefix_sum = 0 은 인덱스 -1 로 간주 (길이 조건을 자연스럽게 맞추기 위함)

        prefix_sum = 0                       # 현재까지의 prefix sum
        max_sum = -inf                       # 최대 subarray 합을 추적

        for i, num in enumerate(nums):
            prefix_sum += num                # prefix_sum = nums[0..i] 의 누적합
            mod = i % k                      # 현재 prefix 인덱스 i에 대한 mod 그룹 결정

            old_pre = least_pre[mod]         # 같은 mod 그룹에서 지금까지 가장 작은 prefix sum

            sub_sum = prefix_sum - old_pre   # 길이가 k의 배수인 subarray 후보 합 = 현재 prefix - 최소 prefix

            if max_sum < sub_sum:            # 최대 합 갱신 여부 확인
                max_sum = sub_sum

            if old_pre > prefix_sum:         # 현재 prefix_sum이 이 mod 그룹의 최소값을 갱신할 수 있으면 업데이트
                least_pre[mod] = prefix_sum

        return max_sum                       # 길이가 k의 배수인 subarray 합 중 최댓값 반환

4. 마무리

혼자 힘으로 풀고 싶다.