최대 하나의 원소를 변경한 후 가장 긴 등차수열 - LeetCode-3872 풀이

Source: Dev.to

Source: …

하나의 원소를 최대 하나까지 바꿨을 때 가장 긴 등차 수열

우리는 정수 배열 nums가 주어집니다.
부분 배열이 등차수열이라는 것은 연속된 원소들 사이의 차이가 일정함을 의미합니다.
우리는 최대 하나의 원소를 원하는 어떤 정수로든 교체할 수 있습니다.

목표: 단 한 번의 교체 후 얻을 수 있는 등차 부분 배열의 최대 길이를 반환합니다.

예시

Input:  nums = [9, 7, 5, 10, 1]
Output: 5

10을 3으로 교체 → [9, 7, 5, 3, 1].
공차 = -2. 길이 = 5.

왜 순진한 O(N²) 해결책은 작동하지 않을까

명백한 접근법은 다음과 같다:

1. 모든 인덱스 i를 순회한다.
2. nums[i]를 주변 요소에 맞는 값으로 교체해 본다.
3. 좌우로 스캔하여 가장 긴 등차 부분 배열을 찾는다.

각 인덱스마다 이를 수행하면 **O(N²)**가 되며, 이는 제한인 N ≤ 10⁵를 초과한다.
O(N) 해결책이 필요하다.

핵심 관찰

우리가 교체하는 원소는 다리 역할을 하여 연결할 수 있습니다:

• 왼쪽에 있는 유효한 등차수열, 그리고
• 오른쪽에 있는 유효한 등차수열.

모든 교체를 시뮬레이션하는 대신, 두 개의 보조 배열을 미리 계산합니다:

두 배열은 한 번의 선형 탐색으로 구축됩니다.

Source: …

left와 right 사용 방법

각 인덱스 i(교체할 수 있는 요소)에 대해

• l = i‑1 (왼쪽 이웃)
• r = i+1 (오른쪽 이웃)

가능한 경우는 세 가지입니다:

1. 왼쪽만 확장 – 교체된 요소를 l에서 끝나는 수열에 추가합니다.
   len = left[l] + 1
2. 오른쪽만 확장 – 교체된 요소를 r에서 시작하는 수열 앞에 삽입합니다.
   len = right[r] + 1
3. 양쪽을 연결 – nums[i]를 nums[l]과 nums[r] 사이의 정확한 중간값으로 만듭니다.
   • 이는 (nums[r] - nums[l])가 짝수일 때만 가능합니다.
   • 필요한 공차: diff = (nums[r] - nums[l]) / 2.
   • 왼쪽과 오른쪽 수열이 모두 이 diff를 갖는다면, i를 통해 두 수열을 합칠 수 있습니다.

연결(bridge) 길이는 3(nums[l], 교체된 nums[i], nums[r])에서 시작하며, 동일한 diff를 공유하는 왼쪽/오른쪽 수열의 길이만큼 확장됩니다.

left와 right 만들기

# left[i] = length of longest arithmetic subarray ending at i
left = [1] * n
for i in range(2, n):
    if nums[i] - nums[i-1] == nums[i-1] - nums[i-2]:
        left[i] = left[i-1] + 1
    else:
        left[i] = 2          # start a new pair

# right[i] = length of longest arithmetic subarray starting at i
right = [1] * n
for i in range(n-3, -1, -1):
    if nums[i+1] - nums[i] == nums[i+2] - nums[i+1]:
        right[i] = right[i+1] + 1
    else:
        right[i] = 2          # start a new pair

두 루프는 O(N) 시간에 실행됩니다.

Source:

전체 구현

파이썬

from typing import List

class Solution:
    def longestArithmetic(self, nums: List[int]) -> int:
        n = len(nums)
        if n < 2:
            return n

        # pre‑compute left and right arrays
        left = [1] * n
        for i in range(2, n):
            if nums[i] - nums[i-1] == nums[i-1] - nums[i-2]:
                left[i] = left[i-1] + 1
            else:
                left[i] = 2

        right = [1] * n
        for i in range(n-3, -1, -1):
            if nums[i+1] - nums[i] == nums[i+2] - nums[i+1]:
                right[i] = right[i+1] + 1
            else:
                right[i] = 2

        ans = 2  # at least two elements form an arithmetic subarray

        # try changing each element
        for i in range(1, n-1):
            # case 1: extend left side only
            ans = max(ans, left[i-1] + 1)
            # case 2: extend right side only
            ans = max(ans, right[i+1] + 1)

            # case 3: bridge both sides
            if (nums[i+1] - nums[i-1]) % 2 == 0:
                diff = (nums[i+1] - nums[i-1]) // 2
                cur_len = 3
                # extend left while maintaining diff
                l = i-2
                while l >= 0 and nums[l+1] - nums[l] == diff:
                    cur_len += 1
                    l -= 1
                # extend right while maintaining diff
                r = i+2
                while r < n and nums[r] - nums[r-1] == diff:
                    cur_len += 1
                    r += 1
                ans = max(ans, cur_len)

        # edge cases: modify first or last element
        ans = max(ans, right[1] + 1)   # change nums[0]
        ans = max(ans, left[n-2] + 1)  # change nums[n-1]

        return ans

C++ (원본 그대로)

// Note: The original snippet is incomplete; it is reproduced verbatim.
if (n  left(n, 2), right(n, 2);
left[0] = 1;               // only the first element
right[n-1] = 1;            // only the last element

// ----- left -----
for (int i = 2; i = 0; --i) {
    if (nums[i+1] - nums[i] == nums[i+2] - nums[i+1])
        right[i] = right[i+1] + 1;
}

int ans = 2;

// Edge cases: modify first or last element
ans = max(ans, right[1] + 1);          // change nums[0]
ans = max(ans, left[n-2] + 1);         // change nums[n-1]

// ----- try modifying each middle element -----
for (int i = 1; i  0 && (nums[l] - nums[l-1]) == diff)
            cur_len += left[l] - 1;      // left[l] already counts nums[l]

        if (r & nums) {

인사이트

Precomputing prefix/suffix states and combining them at each index은 다음과 같은 문제에서 지속적으로 등장합니다:

• 최대 하나의 교체 / 삭제
• 단일 요소를 통해 두 개의 유효한 부분 배열을 병합
• 하나의 “무료 패스”가 있는 슬라이딩‑윈도우 변형

요점

“최대 하나의 요소를 변경한다”는 경우를 볼 때마다, 왼쪽과 오른쪽에서 미리 계산할 수 있는 것이 무엇인지 생각해 보세요.