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

Python으로 정확히 k개의 서로 다른 단어를 포함하는 부분 리스트 개수 구하기


문제 이해하기

단어로 이루어진 리스트와 정수 k가 주어졌을 때, 주어진 리스트에서 정확히 k개의 서로 다른 단어를 포함하는 연속된 부분 리스트(sublist)의 개수를 구하는 문제입니다.

예를 들어, words = ["Kolkata", "Delhi", "Delhi", "Kolkata"]이고 k = 2라고 가정해 보겠습니다. 이때 출력은 5가 됩니다. 다음과 같은 부분 리스트들이 정확히 2개의 고유한 단어를 포함하기 때문입니다.

  • ["Kolkata", "Delhi"]
  • ["Delhi", "Kolkata"]
  • ["Kolkata", "Delhi", "Delhi"]
  • ["Delhi", "Delhi", "Kolkata"]
  • ["Kolkata", "Delhi", "Delhi", "Kolkata"]

반면 ["Delhi", "Delhi"]는 고유한 단어가 하나뿐이므로 개수에 포함되지 않습니다.

해결 접근 방법: 슬라이딩 윈도우

이 문제는 슬라이딩 윈도우(Sliding Window) 기법으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

먼저 '최대 k개의 서로 다른 단어를 포함하는 부분 리스트의 개수'를 세는 함수 work()를 정의합니다. 그러면 '정확히 k개'인 경우의 수는 아래와 같이 계산할 수 있습니다.

work(words, k) - work(words, k - 1)

'최대 k개'에서 '최대 k-1개'를 빼면 자연스럽게 '정확히 k개'인 경우만 남게 됩니다.

work() 함수의 동작 단계

  • work(words, k) 함수를 정의합니다.
  • n := words의 길이로 설정합니다.
  • k가 0이면 0을 반환합니다.
  • cnt := 새로운 딕셔너리(각 단어의 등장 횟수 저장), ans := 0, l := 0(왼쪽 포인터)으로 초기화합니다.
  • r을 0부터 n-1까지 반복합니다.
    • word := words[r]을 꺼내고, cnt에 없으면 cnt[word] := 0으로 초기화한 뒤 cnt[word]를 1 증가시킵니다.
    • cnt의 크기(서로 다른 단어 수)가 k보다 커지면, 왼쪽 끝 단어 words[l]의 개수를 1 줄이고, 0이 되면 딕셔너리에서 제거한 후 l을 1 증가시킵니다.
    • ans := ans + (r - l + 1)을 누적합니다. 이는 오른쪽 끝이 r인 유효한 부분 리스트의 개수입니다.
  • 모든 반복이 끝나면 ans를 반환합니다.

Python 코드 예제

아래 구현을 통해 더 잘 이해해 보겠습니다.

class Solution:
   def solve(self, words, k):
      return self.work(words, k) - self.work(words, k - 1)
   def work(self, words, k):
      n = len(words)
      if k == 0:
         return 0
      cnt = dict()
      ans = 0
      l = 0
      for r in range(n):
         word = words[r]
         if word not in cnt:
            cnt[word] = 0
         cnt[word] += 1
         while len(cnt) > k:
            cnt[words[l]] -= 1
            if cnt[words[l]] == 0:
               del cnt[words[l]]
            l += 1
         ans += r - l + 1
      return ans
ob = Solution()
words = ["Kolkata", "Delhi", "Delhi", "Kolkata"]
k = 2
print(ob.solve(words, k))

입력

["Kolkata", "Delhi", "Delhi", "Kolkata"], 2

출력

5

복잡도 분석

슬라이딩 윈도우 기법 덕분에 시간 복잡도는 O(n)이며, 공간 복잡도는 딕셔너리에 저장되는 서로 다른 단어의 수에 비례하여 O(n)입니다. 브루트포스 방식(O(n²))보다 훨씬 효율적으로 동작하므로, 리스트의 길이가 클 때도 안정적인 성능을 기대할 수 있습니다.