Program to find number of sublists that contains exactly k different words in Python

Suppose we have a list of words and another value k. We have to find the number of sublists in the given words such that there are exactly k different words.

So, if the input is like words = ["Kolkata", "Delhi", "Delhi", "Kolkata"] k = 2, then the output will be 5, as the following sublists have 2 unique words: ["Kolkata", "Delhi"], ["Delhi", "Kolkata"],["Kolkata","Delhi","Delhi"], ["Delhi","Delhi","Kolkata"], ["Kolkata","Delhi","Delhi","Kolkata"], but not ["Delhi","Delhi"] as there is only one unique word.

To solve this, we will follow these steps −

• Define a function work() . This will take words, k
• n := size of words
• if k is same as 0, then
  • return 0
• cnt := a new map
• ans := 0
• l := 0
• for r in range 0 to n, do
  • word := words[r]
  • if word is not present in cnt, then
    • cnt[word] := 0
  • cnt[word] := cnt[word] + 1
  • while size of cnt > k, do
    • cnt[words[l]] := cnt[words[l]] - 1
    • if cnt[words[l]] is same as 0, then
      • remove words[l] from cnt
    • l := l + 1
  • ans := ans + r - l + 1
• return ans

From the main method return (work(words, k) - work(words, k - 1))

Example (Python)

Let us see the following implementation to get better understanding −

 Live Demo

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))

Input

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

Input

5