Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 홀수·짝수 인덱스 요소의 합을 같게 만드는 제거 가능 요소 개수 찾기

문제 설명

숫자로 이루어진 리스트 nums가 주어져 있다고 가정해 봅시다. 함수 f(i)는 인덱스 i에 위치한 요소를 삭제한 후, 결과 리스트에서 짝수 인덱스 요소들의 합홀수 인덱스 요소들의 합이 서로 같으면 참(True), 그렇지 않으면 거짓(False)을 반환합니다. 이때 우리가 구해야 하는 것은 함수 f가 참을 반환하도록 만드는 인덱스의 총개수입니다.

예를 들어 입력이 nums = [6, 8, 5, 2, 3]이라면 출력은 2입니다. 그 이유는 다음과 같습니다.

  • 8을 제거하면 리스트는 [6, 5, 2, 3]이 되고, 짝수·홀수 인덱스 요소들의 합은 각각 8로 동일합니다.
  • 2를 제거하면 리스트는 [6, 8, 5, 3]이 되고, 짝수·홀수 인덱스 요소들의 합은 각각 11로 동일합니다.

접근 방법

매번 요소를 하나씩 실제로 제거하고 전체 합을 다시 계산하는 비효율적인 방법 대신, 누적 합(prefix sum)을 미리 구해 두면 각 인덱스를 상수 시간에 검사할 수 있어 전체 문제를 O(n) 시간 복잡도로 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 인덱스 i의 요소를 제거하면, i 앞쪽 요소들은 기존 인덱스의 홀짝 성격을 그대로 유지하는 반면, i 뒤쪽 요소들은 인덱스가 하나씩 당겨지면서 홀짝이 반전됩니다. 따라서 제거 후 짝수 인덱스 그룹의 합은 'i 이전의 짝수 인덱스 합'과 'i 이후의 홀수 인덱스 합'을 더한 값이 됩니다. 이 값이 남은 전체 합의 절반과 일치하면, 즉 e * 2 == s - nums[i]를 만족하면 두 그룹의 합이 같아집니다.

단계별 풀이

  • n := nums의 길이
  • a := 2 × (n+1) 크기의 2차원 리스트를 생성하고 모든 값을 0으로 초기화 (0행은 짝수 인덱스, 1행은 홀수 인덱스의 누적 합 저장)
  • nums의 각 인덱스 i와 값 x에 대해 다음을 반복합니다.
    • a[0][i + 1] := a[0][i]
    • a[1][i + 1] := a[1][i]
    • a[i mod 2][i + 1] := a[i mod 2][i + 1] + x → 현재 인덱스의 홀짝에 맞는 행에 값 누적
  • c := 0 (조건을 만족하는 인덱스 개수)
  • s := nums의 전체 요소 합
  • i를 0부터 n-1까지 순회하며 다음을 수행합니다.
    • e := a[0][i] - a[0][0] + a[1][n] - a[1][i + 1] → i번째 요소 제거 후 짝수 인덱스 그룹의 합
    • 만약 e * 2 == s - nums[i]이면 (홀수 인덱스 그룹의 합과 같아지는 조건)
      • c := c + 1
  • c 반환

구현 예제

다음 파이썬 구현을 통해 더 자세히 이해해 보겠습니다.

def solve(nums):
    n = len(nums)
    a = [[0] * (n + 1), [0] * (n + 1)]
    for i, x in enumerate(nums):
        a[0][i + 1] = a[0][i]
        a[1][i + 1] = a[1][i]
        a[i % 2][i + 1] += x

    c = 0
    s = sum(nums)
    for i in range(n):
        e = a[0][i] - a[0][0] + a[1][n] - a[1][i + 1]
        if e * 2 == s - nums[i]:
            c += 1
    return c

nums = [6, 8, 5, 2, 3]
print(solve(nums))

입력

[6, 8, 5, 2, 3]

출력

2