Python으로 첫 번째 요소와 마지막 요소가 같은 부분 리스트 개수 구하기

숫자로 이루어진 리스트 nums가 주어졌을 때, 첫 번째 요소와 마지막 요소가 동일한 부분 리스트(sublist)의 개수를 구하는 문제입니다.

예를 들어 입력이 nums = [10, 15, 13, 10]이라면 결과는 5가 됩니다. 조건을 만족하는 부분 리스트는 다음과 같습니다.

• [10]
• [15]
• [13]
• [10]
• [10, 15, 13, 10]

접근 방법

길이가 1인 부분 리스트는 항상 첫 번째 요소와 마지막 요소가 같으므로, 먼저 전체 리스트 길이만큼 개수를 초기화합니다. 이후 각 숫자의 등장 빈도를 세고, 같은 숫자가 v번 등장한다면 그중 두 개를 뽑아 시작과 끝으로 삼는 경우의 수인 v × (v − 1) / 2만큼 추가하면 됩니다. 이는 조합 공식 C(v, 2)에 해당합니다.

알고리즘 단계

1. num_sublists := len(nums) 로 초기화합니다.
2. 빈 딕셔너리(맵) d를 준비합니다.
3. nums의 각 원소 n에 대해 d[n] 값을 1씩 증가시켜 빈도를 계산합니다.
4. d의 각 숫자 k와 빈도 v에 대해, v가 1이 아니라면 num_sublists에 (v−1) × v / 2를 더합니다.
5. num_sublists를 반환합니다.

이 방식은 모든 부분 리스트를 직접 탐색하는 O(n²) 대신 O(n) 시간 복잡도로 문제를 해결할 수 있다는 장점이 있습니다.

구현 예제

from collections import defaultdict
class Solution:
    def solve(self, nums):
        num_sublists = len(nums)

        d = defaultdict(int)
        for n in nums:
            d[n] += 1

        for k, v in d.items():
            if v != 1:
                num_sublists += (v - 1) * (v) // 2
        return num_sublists

ob = Solution()
nums = [10, 15, 13, 10]
print(ob.solve(nums))

입력

[10, 15, 13, 10]

출력

5

위 코드에서 defaultdict(int)를 사용하면 처음 등장하는 숫자에 대해 KeyError 없이 자동으로 0으로 초기화되므로 빈도 계산이 간결해집니다. 정수 나눗셈 연산자 //를 사용해 결과가 항상 정수로 유지됩니다.