Сколько подстрок в строке имеют равное количество 3х букв

Question

Дана строка, состоящая только из букв A, B и C. Необходимо посчитать за O(n), сколько подстрок содержат одинаковое количество всех трех букв.

Например:

CCAABB - 1(CCAABB),

CBACBA - 5(CBA,BAC,ACB,CBA,CCAABB),

AAABBC - 0,

CCBCCCABB - 1(CAB),

CBACBACBA - 12

Есть попытка решить таким образом, но это решение работает только для некоторых случаев. Подскажите, как решить нормально? Рекомендовано использовать сеты или словари, но это не обязательно

def count(s):
  n = len(s)
  result = 0
  count = {}
    
  for i in range(n):
    print(s[i])
    if s[i] not in count:
      count[s[i]] = 1
    else:
      count[s[i]] += 1
    print(count)

    if len(count) == 3 and count["A"] == count["B"] == count["C"]:
      result += (i+1)/3
      
    elif len(count) == 3:
      if abs(count["C"]-count["B"]==1) or abs(count["C"]-count["A"]==1):
        result += i//3
  
  return result

Answer 1

Вот решение в лоб:

text = "CBACBA"

count = 0

for size in range(1, len(text) // 3 + 1):
    for pos in range(0, len(text) + 1 - 3 * size):
        count_a = 0
        count_b = 0
        count_c = 0

        for index in range(3 * size):
            if text[pos + index] == 'A':
                count_a += 1

            if text[pos + index] == 'B':
                count_b += 1

            if text[pos + index] == 'C':
                count_c += 1

        if (count_a == count_b == count_c):
            count = count + 1

print(count)

Правда я бы не назвал бы это O(n), скорее O(n(n - 1) / 6) ну т.е. O(n^2) :)

Измененный алгоритм для честных n^2 (плавающее окно):

text = "CBACBA"

count = 0

for step in range(1, len(text) // 3 + 1):
    count_a = 0
    count_b = 0
    count_c = 0

    size = step * 3

    for pos in range(size):
        if text[pos] == 'A':
            count_a = count_a + 1

        if text[pos] == 'B':
            count_b = count_b + 1

        if text[pos] == 'C':
            count_c = count_c + 1

    if (count_a == count_b == count_c):
        count = count + 1

    for pos in range(1, len(text) + 1 - size):
        if text[pos - 1] == 'A':
            count_a = count_a - 1

        if text[pos - 1] == 'B':
            count_b = count_b - 1

        if text[pos - 1] == 'C':
            count_c = count_c - 1

        if text[pos + size - 1] == 'A':
            count_a = count_a + 1

        if text[pos + size - 1] == 'B':
            count_b = count_b + 1

        if text[pos + size - 1] == 'C':
            count_c = count_c + 1

        if (count_a == count_b == count_c):
            count = count + 1

print(count)

По совету @Danis включил словари для уменьшения кода:

text = "CBACBA"

total_count = 0

for step in range(1, len(text) // 3 + 1):
    local_count = {"A": 0, "B": 0, "C": 0}

    size = step * 3

    for pos in range(size):
        local_count[text[pos]] += 1

    if (local_count['A'] == local_count['B'] == local_count['C']):
        total_count += 1

    for pos in range(1, len(text) + 1 - size):
        local_count[text[pos - 1]] -= 1
        local_count[text[pos + size - 1]] += 1

        if (local_count['A'] == local_count['B'] == local_count['C']):
            total_count += 1

print(total_count)

Answer 2

Просто альтернативный вариант квадратичного алгоритма. Сначала за линейнок время готовим счетчики, потом для каждого интервала (start, end) за константу находим ответ. Проблема только в том, что таких интервалов надо проверить n^2.

Код на C#

int count(string str){
    int ret = 0;
    
    int[] acount = new int[str.Length];
    int[] bcount = new int[str.Length];
    int[] ccount = new int[str.Length];
    
    if (str[0] == 'A') acount[0] = 1;
    if (str[0] == 'B') bcount[0] = 1;
    if (str[0] == 'C') ccount[0] = 1;

    for (int i = 1; i < str.Length; i++)
    {
        acount[i] = acount[i-1];
        bcount[i] = bcount[i-1];
        ccount[i] = ccount[i-1];
        
        if (str[i] == 'A') acount[i]++;     
        if (str[i] == 'B') bcount[i]++; 
        if (str[i] == 'C') ccount[i]++; 
    }   
    
    for(int start = 0; start <str.Length-2; start ++){
        for(int end = start +2; end <str.Length; end +=3)
        {
            int a = acount[end] - (start == 0 ? 0 : acount[start-1]);
            int b = bcount[end] - (start == 0 ? 0 : bcount[start-1]);
            int c = ccount[end] - (start == 0 ? 0 : ccount[start-1]);           
            if (a == b && a == c) ret ++;
        }
    }
    
    return ret;
}

Проверка

Console.WriteLine(count("CCAABB"));
Console.WriteLine(count("CBACBA"));
Console.WriteLine(count("AAABBC"));
Console.WriteLine(count("CCBCCCABB"));
Console.WriteLine(count("CBACBACBA"));

Вывод

Answer 3

Линейный. Составной ключ diff соответствует текущему соотношению счетчиков A к B и А к С

Ключ - относительное соотношение счётчиков для каждого из символов. Например, (-1,2) означает, что до текущей позиции количество B на 1 меньше, чем A, а С - на 2 больше, чем A. Если это соотношение совпадает для каких-то позиций, то между ними - одинаковое количество всех символов. Например, указанный ключ будет для BACACCC и BACACCCCACABB, значит суффикс CACABB содержит поровну всех символов. Такой же ключ будет для BACACCCCACABBBCA, но на этот момент его счётчик уже равен двум, так что есть два суффикса, кончающихся в данной позиции, и содержащие поровну символов.

https://ideone.com/MOIMqj

from collections import Counter
def count(s):
    dic = Counter()
    diff = (0,0)
    dic[diff] += 1
    result = 0
    for c in s:
        if c == 'A':
            diff = (diff[0]-1,diff[1]-1)
        elif c == 'B':
            diff = (diff[0]+1,diff[1])
        else:
            diff = (diff[0],diff[1]+1)
        result += dic[diff]
        dic[diff] += 1
    return result


print(count('CBACBACBA'))

БЛОГ НА HUSL

Сколько подстрок в строке имеют равное количество 3х букв

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