Посчитать сколько подстрок в строке имеют именно 2 различные буквы

Question

Необходимо посчитать сколько вариантов подстроки включает в себя строка, чтобы были именно 2 буквы уникальные. То есть "aabacba" будет сожержать в себе 8 (aab, aaba, ab, aba, ac, ba, ba, cb) "abab" - 6(ab, ba, aba, ab, bab, abab), а "aaaa" ни одной.

Два дня голову ломаю. Получился такой вариант, но надо, чтобы это было О(n), иначе не проходит временные тесты. То есть по всей видимости через for надо и без вложенных циклов. Или я что-то упускаю? Подскажите, пожалуйста.

def count(s):
  sum = 0
  n = len(s)
  i = 0
  index = 0
  charList = [0] * 27
  dist1 = ""
  dist2 = ""
    
  while i < n:
    if (index < n) and (s[index]==dist1 or s[index]==dist2 or dist1=="" or dist2==""):
      if charList[ord(s[index]) - 97] == 0 and charList[ord(s[index-1]) - 97] != 0 and dist2 == "":
        dist2 = s[index]
      elif charList[ord(s[index]) - 97] == 0 and dist1 == "":
        dist1 = s[index]
 
      if s[index]==dist1 or s[index]==dist2: 
        charList[ord(s[index]) - 97] += 1 
        if dist1 != "" and dist2 != "": 
          sum += 1
      index += 1    
    else:
      dist1 = ""
      dist2 = ""
      charList = [0] * 27
      i += 1
      index = i
    
      return sum

if __name__ == "__main__":
    print(count("babbbabbba")) # 39
    print(count("aaaa")) # 0
    print(count("abab")) # 6 
    print(count("aaba")) #5 
    print(count("abbba")) #7 
    print(count("aabacba")) # 8 
    print(count("ababa")) # 10

Answer 1

from itertools import groupby


def count(example):

    answer, sum_n_char = 0, 0
    curr_chars = set()
    prev_char = None

    for char, n_char in groupby(example):
        n_char = sum(1 for elem in n_char)
        curr_chars.add(char)

        if len(curr_chars) == 2:
            answer += sum_n_char * n_char

        elif len(curr_chars) > 2:
            answer += prev_n_char * n_char
            curr_chars = set()
            curr_chars.add(char)
            sum_n_char = 0

        prev_char, prev_n_char = char, n_char
        sum_n_char += n_char

    return answer

Краткая идея:

Делаем RLE преобразование, т.е. кодируем строку из aabbc -> 2a2b1c
Далее вариант такой, либо до этого у нас было 2 символа и на текущем шаге берем тот же самый из этих двух, то мы можем выбрать любую первую позицию из первой группы для старта подстроки и любую конечную из второй. (для примера, если было 2a2b2a и новая группа 2b, то у нас 6 * 2 варианта, т.к. из первой группы для начала 6 вариантов и из второй группы для конца 2 варианта).
Если мы встретили новый символ и он нам дал 3-ий символ, то мы все пред. группы забываем, кроме последней и ведем отсчет с нее.

Описание не совсем точное, но если понять суть решения, то вы сможете воспроизвести недостающие части.

БЛОГ НА HUSL

Посчитать сколько подстрок в строке имеют именно 2 различные буквы

Ответы (1 шт):