1. 문제 소개

3652. Best Time to Buy and Sell Stock using Strategy

2. 나의 풀이

처음에는 모든 합을 구한 다음에, 변화량만 추적했었다.
코드 자체는 맞게 풀었지만 시간초과로 인해 틀리고 말았다.

class Solution:
    def maxProfit(self, prices: List[int], strategy: List[int], k: int) -> int:
        # 초기 전략(strategy)을 그대로 적용했을 때의 기본 수익
        # strategy[i] ==  1 : +prices[i]
        # strategy[i] ==  0 : 0
        # strategy[i] == -1 : -prices[i]
        base = sum(p * s for p, s in zip(prices, strategy))
        n = len(prices)

        ans = base  # 최대 수익 초기값은 기본 수익

        # 길이 k인 윈도우를 왼쪽부터 하나씩 이동
        for i in range(n - k + 1):
            delta = 0  # 현재 윈도우에서 전략 변경으로 인해 발생하는 수익 변화량

            # [앞 절반] i ~ i + k//2 - 1
            # 전략을 모두 0으로 변경한다고 가정
            for idx in range(i, i + k // 2):
                if strategy[idx] == 1:
                    # 기존 +prices[idx] → 0 이 되므로 수익 감소
                    delta -= prices[idx]
                elif strategy[idx] == -1:
                    # 기존 -prices[idx] → 0 이 되므로 수익 증가
                    delta += prices[idx]

            # [뒤 절반] i + k//2 ~ i + k - 1
            # 전략을 모두 +1로 변경한다고 가정
            for idx in range(i + k // 2, i + k):
                if strategy[idx] == 0:
                    # 기존 0 → +prices[idx]
                    delta += prices[idx]
                elif strategy[idx] == -1:
                    # 기존 -prices[idx] → +prices[idx]
                    # 총 변화량은 +2 * prices[idx]
                    delta += 2 * prices[idx]

            # 현재 윈도우에서의 총 수익 = base + delta
            # 최대값 갱신
            if base + delta > ans:
                ans = base + delta

        return ans  # 가능한 최대 수익 반환

3. 다른 풀이

결국 더 prefix + sliding window 2가지 방식 모두 적용한 풀이로 가능했다.
시간복잡도는 $O(n)$이다.

class Solution:
    def maxProfit(self, prices: List[int], strategy: List[int], k: int) -> int:
        ps = [p * s for p, s in zip(prices, strategy)]  # 기본 전략 기여값 배열 ps[i] = prices[i] * strategy[i]
        n = len(prices)                                 # 전체 길이

        base = sum(ps)                                  # 기본 수익(전략 변경 전 총합)
        h = k // 2                                      # 구간을 전/후반으로 나누는 기준 (앞 절반 길이)

        old = sum(ps[:k])                               # 초기 윈도우(0..k-1)에서의 "기존 전략 기여" 합
        new = sum(prices[h:k])                          # 초기 윈도우에서 "변경 후 기여"로 계산되는 부분 합
                                                       # (구간 후반 [h..k-1]을 +1로 본 기여로 해석되는 합)

        diff = max(0, new - old)                        # 첫 윈도우에서의 이득(변화량) 후보 (이득이 음수면 0)

        # 오른쪽 끝을 k부터 n-1까지 이동시키며 길이 k 슬라이딩 윈도우 수행
        for right in range(k, n):
            left = right - k + 1                        # 현재 윈도우의 왼쪽 인덱스

            # old 갱신:
            # - 윈도우에 새로 들어온 ps[right]를 더하고
            # - 윈도우에서 빠져나간 ps[left-1]를 뺀다
            old += ps[right] - ps[left - 1]

            # new 갱신:
            # - 변경 후 기여에서 새로 포함되는 prices[right]를 더한다
            # - 변경 후 기여에서 제외되어야 하는 prices[left - 1 + h]를 뺀다
            #   (윈도우가 한 칸 이동하면 후반 구간도 함께 한 칸 이동하므로, 그에 맞춰 빠지는 항을 제거)
            new += prices[right]
            new -= prices[left - 1 + h]

            # 현재 윈도우에서의 이득(new-old)이 지금까지의 최대 diff보다 크면 갱신
            if new - old > diff:
                diff = new - old

        return base + diff                               # 기본 수익 + 최대 변화량(이득) 반환

4. 마무리

아쉬웠지만 그래도 잘 풀었다.