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

Python으로 배열에서 x+1이 함께 존재하는 요소의 개수 구하기

숫자로 이루어진 리스트 nums가 주어졌다고 가정해 봅시다. 이 문제의 목표는 배열 안에서 x + 1 역시 배열에 존재하는 요소 x의 개수를 구하는 것입니다.

예를 들어 입력이 nums = [4, 2, 3, 3, 7, 9]라고 한다면 결과는 3이 됩니다. 2+1=3이 배열에 존재하고, 3+1=4 역시 존재하며, 중복된 값 3 하나까지 포함되어 총 3개이기 때문입니다.

문제 해결 접근 방법

이 문제는 Counter를 활용하면 효율적으로 해결할 수 있습니다. 다음과 같은 단계를 따릅니다.

  • answer를 0으로 초기화합니다.

  • c := nums에 있는 각 요소의 빈도수를 저장한 Counter 객체를 생성합니다.

  • dlist := c의 모든 키(key)로 구성된 리스트를 만듭니다.

  • dlist의 각 요소 i에 대해 다음을 수행합니다.

    • 만약 c[i + 1] > 0, 즉 i+1이 배열에 존재한다면,

      • answerc[i](요소 i의 빈도수)를 더합니다.

  • 모든 반복이 끝나면 최종적으로 answer를 반환합니다.

이 방식은 배열을 한 번 순회하여 빈도수를 계산한 뒤(O(n)), 고유한 키의 개수(k)만큼만 반복하기 때문에 전체 시간 복잡도는 O(n)으로 매우 효율적입니다.

구현 예제

아래 구현 예제를 통해 더 자세히 이해해 보겠습니다.

from collections import Counter
def solve(nums):
   answer = 0
   c = Counter(nums)
   dlist = list(c.keys())
   for i in dlist:
      if c[i + 1] > 0:
         answer += c[i]

   return answer

nums = [4, 2, 3, 3, 7, 9]
print(solve(nums))

입력

[4, 2, 3, 3, 7, 9]

출력

3