Правильные подматрицы задачи

Ниже приведено одно из возможных решений на Python с подробными комментариями. Общая идея такова:

1. Разделяем рассмотрение на два случая:

   • Подматрицы высотой в ровно 1 строку.
   • Подматрицы высотой не менее 2 строк.

2. Случай высотой 1 (одна строка):

   • В одной строке цвета (чёрные/белые клетки) идут вперемешку: 0,1,0,1,... (если считать $(r + c) \bmod 2$ ).
   • Чтобы подотрезок (в одной строке это "подматрица") был «правильным», он обязан содержать хотя бы одну чёрную и одну белую клетку, и при этом все белые клетки в нём должны содержать одну и ту же букву $p$, а все чёрные — одну и ту же букву $q$, причём $p \neq q$.
   • Для фиксированной строки мы идём слева направо, «накапливая» текущие $p$ и $q$. Если встречаем клетку белого цвета, которая не совпадает с уже зафиксированным $p$ (или если $p$ ещё не зафиксирован, тогда берём его равным этой букве), то текущее "окно" заканчиваем и начинаем новое. Аналогично для чёрной клетки и $q$.
   • В итоге мы получим участки (колонки подряд) длины $L$, внутри которых белые клетки содержат ровно $p$, а чёрные — ровно $q$. Если в таком участке действительно встречались и белые, и чёрные клетки (то есть строка не «однородна» по цветам внутри этого участка), и $p\neq q$, то все подотрезки длиной $\ge 2$ внутри этого участка будут «правильными». Количество таких подотрезков внутри куска длины $L$ равно $$\text{(все подотрезки)} - \text{(одиночные)} \;=\; \frac{L(L+1)}{2} \;-\; L \;=\;\frac{L(L-1)}{2}.$$
   • Суммируем по всем подобным участкам для каждой строки.

3. Случай высотой $\ge 2$:

   • Если подматрица покрывает хотя бы 2 строки, то в ней гарантированно есть и белые, и чёрные клетки (при условии, что она содержит хотя бы 1 столбец). Поэтому отдельное требование «у подматрицы должна быть хотя бы одна белая и одна чёрная клетка» автоматически выполняется, если высота $\ge 2$.
   • Чтобы подматрица была «правильной», все белые клетки внутри неё должны содержать одну и ту же букву $p$, а все чёрные — одну и ту же букву $q$, причём $p \neq q$.

   Можно делать так:

   1. Предварительно посчитаем по строкам префиксные частоты для каждого столбца и каждой буквы — отдельно для белых клеток (color=0) и для чёрных (color=1). То есть заведём структуры вида

      $$\text{freq0}[r][c][\text{letter}],\quad \text{freq1}[r][c][\text{letter}],$$

      где $\text{freq0}[r][c][\ell]$ означает: «Сколько раз буква $\ell$ встретилась среди белых клеток в столбце $c$ от строки $0$ до строки $r-1$ включительно». Аналогично $\text{freq1}$ для чёрных клеток.

   2. Перебираем все пары строк $(r_1, r_2)$ с $0 \le r_1 < r_2 < n$. (Высота $\ge 2$.)

   3. Для каждого столбца $c$ вычисляем, единственная ли буква встречается на белых клетках в диапазоне строк $[r_1..r_2]$.

      • Для этого найдём, сколько всего белых клеток в столбце $c$ от $r_1$ до $r_2$. Это: $$\text{count0} = \text{(число белых клеток)} = \text{freq0\_ALL}(r_2+1, c) - \text{freq0\_ALL}(r_1, c),$$ где $\text{freq0\_ALL}$ — это сумма по всем буквам (просто счётчик, сколько белых клеток всего).
      • Затем смотрим для каждой буквы $\ell$ её частоту $\text{cur} = \text{freq0}[r_2+1][c][\ell] - \text{freq0}[r_1][c][\ell]$. Если ровно одна буква $\ell$ даёт $\text{cur} = \text{count0}$, а у остальных 0, то белые клетки однородны и равны этой букве $\ell$. Если такой $\ell$ нет (или нашлось больше одной подходящей) — значит столбец $c$ «не годится» для построения правильной подматрицы по белым клеткам в этом вертикальном срезе. То же делаем и для чёрных клеток.

   4. Таким образом для каждого столбца $c$ получаем пару $(p_c,\;q_c)$, где $p_c$ — буква, единственная на белых клетках (или $-1$, если не однородно), $q_c$ — буква, единственная на чёрных клетках (или $-1$, если не однородно).

      • Подматрица будет правильной на строках $[r_1..r_2]$, если мы возьмём непрерывный блок столбцов $c_1..c_2$ и у всех этих столбцов будет одна и та же пара $(p,q)$ с $p \neq -1, q \neq -1, p \neq q$.

   5. Следовательно, мы просто идём по массиву столбцов $\text{colPair}[c] = (p_c, q_c)$, группируя одинаковые подряд и проверяя $p_c \neq q_c$. Если длина группы столбцов равна $L$, то количество способов выбрать подотрезок (по столбцам) есть $\tfrac{L(L+1)}{2}$.

   В сумме это даёт количество «правильных» подматриц высотой $\ge 2$.

4. Складываем результаты от (1) и (2).

Важно: в случае высоты $\ge 2$ любая выбранная полоса столбцов уже содержит и белые, и чёрные клетки (по вертикали там точно два разных цвета). Поэтому дополнительное условие «должны быть и белые, и чёрные» автоматически выполнится, и вычитать ничего не надо.

5. Проверяем отдельно, что одиночные клетки не считаются: по условию нужно ровно две разные буквы в шахматном порядке, значит размер подматрицы должен быть $\ge 2$ клетки. В нашем алгоритме:
   • В случае одной строки мы явно отсекаем подотрезки длиной 1 (см. $\tfrac{L(L-1)}{2}$).
   • В случае двух и более строк одиночных подматриц по вертикали не образуется (минимум 2 клетки по вертикали).

Ниже код, который реализует эту идею. Он достаточно большой по количеству операций, но при аккуратной реализации на C++ часто может уложиться в стандартные лимиты при $n,m \le 300$. На Python подобная реализация может потребовать предельно возможной оптимизации (или может оказаться слишком медленной в зависимости от жёсткости лимитов).

python
def solve():
    import sys
    data = sys.stdin.read().strip().split()
    n, m = map(int, data[:2])
    matrix_strings = data[2:]
    
    # Прочитаем матрицу как список строк:
    matrix = [list(row) for row in matrix_strings]
    
    # Преобразуем буквы 'a'..'z' в числа 0..25
    mat = [[ord(matrix[r][c]) - ord('a') for c in range(m)] for r in range(n)]
    
    # color0[r][c] = True, если (r + c) % 2 == 0 (белая клетка), иначе False
    # color1[r][c] = нечто обратное
    # Чтобы не хранить отдельно, просто будем проверять через (r+c)%2
    def is_color0(r, c):
        return ((r + c) & 1) == 0
    
    # ---------------------------------------------------------------------
    # 1) Подсчёт результата для подматриц высотой ровно 1 (каждая строка)
    # ---------------------------------------------------------------------
    answer = 0
    
    for r in range(n):
        # Пройдём по столбцам и будем "накапливать" текущий однородный сегмент
        # p = буква для белых клеток в этом сегменте, q = для чёрных
        # Если встречаем противоречие, закрываем предыдущий сегмент, начинаем новый.
        
        current_p = -1   # какая буква на белых клетках в текущем сегменте
        current_q = -1   # какая буква на чёрных клетках в текущем сегменте
        start_col = 0    # начало текущего сегмента
        have_white = False
        have_black = False
        
        def close_segment(seg_len):
            # seg_len = длина сегмента столбцов подряд
            # Если сегмент "валиден" (have_white and have_black and p!=q),
            # добавляем кол-во подотрезков длины >=2: L*(L-1)/2.
            if seg_len <= 0:
                return 0
            return seg_len * (seg_len - 1) // 2
        
        total_in_this_row = 0
        
        for c in range(m):
            letter = mat[r][c]
            c_color0 = is_color0(r, c)  # True/False
            
            if c_color0:
                # белая клетка => должна совпадать с current_p (если уже зафиксирован)
                # или зафиксировать (если -1)
                if current_p == -1:
                    # ещё не было белых клеток в этом сегменте
                    current_p = letter
                    have_white = True
                elif current_p != letter:
                    # сегмент ломается, нужно закрыть
                    seg_len = c - start_col
                    total_in_this_row += close_segment(seg_len)
                    # начинаем новый сегмент с текущей колонки
                    current_p = letter
                    current_q = -1  # сбросим, т.к. заново ищем q
                    have_white = True
                    have_black = False
                    start_col = c
            else:
                # чёрная клетка => должна совпадать с current_q
                if current_q == -1:
                    current_q = letter
                    have_black = True
                elif current_q != letter:
                    # сегмент ломается
                    seg_len = c - start_col
                    # проверим, валиден ли прошлый
                    if have_white and have_black and current_p != current_q:
                        total_in_this_row += close_segment(seg_len)
                    # начать новый сегмент
                    current_p = -1
                    current_q = letter
                    have_white = False
                    have_black = True
                    start_col = c
            
            # Обновим флаги, что в текущем сегменте есть белые/чёрные
            if c_color0:
                have_white = True
            else:
                have_black = True
            
            # Также если мы видим, что current_p != -1 и current_q != -1,
            # но вдруг они совпали, сегмент тоже ломается
            # (хотя формально можно было ждать следующую клетку, но
            #  удобнее сразу обрубить, так как p=q недопустимо)
            if current_p != -1 and current_q != -1 and current_p == current_q:
                seg_len = c - start_col
                if have_white and have_black and current_p != current_q:
                    total_in_this_row += close_segment(seg_len)
                # начать новый сегмент
                current_p = -1
                current_q = -1
                have_white = c_color0
                have_black = (not c_color0)
                start_col = c
        
        # Закрыть последний сегмент в конце строки
        seg_len = m - start_col
        if have_white and have_black and current_p != -1 and current_q != -1 and current_p != current_q:
            total_in_this_row += close_segment(seg_len)
        
        answer += total_in_this_row
    
    # ---------------------------------------------------------------------
    # 2) Подсчёт результата для подматриц высотой >= 2
    # ---------------------------------------------------------------------
    # Построим freq0 и freq1:
    # freq0[r][c][letter] = сколько раз буква 'letter' встретилась
    # среди белых клеток в столбце c на строках [0..r-1].
    # Аналогично freq1.
    #
    # Чтобы быстро вычитать отрезок [r1..r2], будем брать:
    #   freq0[r2+1][c][letter] - freq0[r1][c][letter].
    
    # И отдельно храним freq0_ALL[r][c] = сколько белых клеток вообще
    # (любой буквы) на строках [0..r-1] в столбце c.
    # Аналогично freq1_ALL.
    
    # Инициализируем массивы нулями:
    freq0 = [[[0]*26 for _ in range(m)] for __ in range(n+1)]
    freq1 = [[[0]*26 for _ in range(m)] for __ in range(n+1)]
    
    freq0_ALL = [[0]*m for _ in range(n+1)]
    freq1_ALL = [[0]*m for _ in range(n+1)]
    
    # Заполним префиксные массивы
    for r in range(n):
        for c in range(m):
            # копируем сверху
            for let in range(26):
                freq0[r+1][c][let] = freq0[r][c][let]
                freq1[r+1][c][let] = freq1[r][c][let]
            freq0_ALL[r+1][c] = freq0_ALL[r][c]
            freq1_ALL[r+1][c] = freq1_ALL[r][c]
            
            # теперь учтём текущую клетку
            letter = mat[r][c]
            if is_color0(r,c):
                freq0[r+1][c][letter] += 1
                freq0_ALL[r+1][c] += 1
            else:
                freq1[r+1][c][letter] += 1
                freq1_ALL[r+1][c] += 1
    
    # Переберём пары строк (r1, r2) с r1 < r2
    for r1 in range(n-1):
        for r2 in range(r1+1, n):
            # Для каждого столбца c найдём пару (p, q):
            #   p = буква белых клеток в [r1..r2], если все белые там одинаковые, иначе -1
            #   q = буква чёрных клеток в [r1..r2], если все чёрные там одинаковые, иначе -1
            colPairs = []
            for c in range(m):
                # Сначала белые:
                cnt0 = freq0_ALL[r2+1][c] - freq0_ALL[r1][c]
                # если cnt0 == 0 => белые клетки отсутствуют в этом столбце на данном вертикальном отрезке
                # тогда для "шахматного" узора это означает "однородность" по сути с любым p,
                # но по условию задачи надо, чтобы субматрица всё-таки включала хоть один белый?
                # Здесь проще считать, что если нет белых клеток, то этот столбец всё равно "не годится",
                # так как в высоте >= 2 белые клетки точно будут (ведь ширина >= 1, значит в этом столбце
                # есть как минимум две клетки, одна из них белая).
                # Однако действительно ли может случиться cnt0 = 0? Да, если высота = r2-r1+1 = 2,
                # но обе клетки оказались чёрными (например, (r1,c) и (r2,c) обе чёрные). 
                # Тогда "по белым клеткам" он формально однороден (там нет белых, нет противоречий).
                # Но такая подматрица не содержит белых клеток вообще, что противоречит условию
                # (надо иметь и белые, и чёрные). Значит такой столбец целиком "не подходит".
                if cnt0 == 0:
                    p = -1
                else:
                    # Проверим, ровно одна ли буква имеет freq = cnt0
                    # freq0[r2+1][c][L] - freq0[r1][c][L] == cnt0 ?
                    # и у остальных 0 ?
                    # Ускоренный путь: попробуем найти именно ту букву L, у которой freq == cnt0.
                    # Если будет несколько таких, либо не будет, => -1
                    p_found = -1
                    for L in range(26):
                        cur = freq0[r2+1][c][L] - freq0[r1][c][L]
                        if cur == cnt0:
                            if p_found == -1:
                                p_found = L
                            else:
                                # Нашли вторую букву, которая тоже "заполняет" весь столбец
                                p_found = -1
                                break
                    p = p_found
                
                # Аналогично для чёрных
                cnt1 = freq1_ALL[r2+1][c] - freq1_ALL[r1][c]
                if cnt1 == 0:
                    q = -1
                else:
                    q_found = -1
                    for L in range(26):
                        cur = freq1[r2+1][c][L] - freq1[r1][c][L]
                        if cur == cnt1:
                            if q_found == -1:
                                q_found = L
                            else:
                                # нашлась вторая буква => неоднородно
                                q_found = -1
                                break
                    q = q_found
                
                if p != -1 and q != -1 and p != q:
                    colPairs.append((p, q))
                else:
                    colPairs.append((-1, -1))
            
            # Теперь идём по colPairs и группируем подряд идущие одинаковые (p,q) != (-1,-1)
            length_block = 0
            cur_pair = (-1, -1)
            
            for c in range(m):
                if colPairs[c] == cur_pair and cur_pair != (-1, -1):
                    # продолжаем блок
                    length_block += 1
                else:
                    # закрыть предыдущий блок
                    if cur_pair != (-1, -1):
                        # Добавим количество способов выбрать подотрезок столбцов:
                        # L*(L+1)//2
                        answer += length_block*(length_block+1)//2
                    # начать новый блок
                    cur_pair = colPairs[c]
                    if cur_pair != (-1, -1):
                        length_block = 1
                    else:
                        length_block = 0
            
            # закрыть последний блок
            if cur_pair != (-1, -1):
                answer += length_block*(length_block+1)//2
    
    print(answer)

Пояснения по ключевым моментам

1. Обработка одного ряда

   • Мы идём по столбцам в строке $r$. Поскольку в одном ряду цвета клеток чередуются по столбцам (0,1,0,1,...), чтобы подотрезок был «шахматным узором» с двумя буквами $p$ и $q$, необходимо в этом подотрезке иметь хотя бы один столбец цвета 0 и один столбец цвета 1. Причём во всех столбцах цвета 0 должна быть одна и та же буква $p$, а во всех столбцах цвета 1 — одна и та же буква $q$, причём $p \neq q$.
   • Мы «накапливаем» такой сегмент подряд. Если встречается противоречие (другая буква в нужном цвете или $p=q$), завершаем сегмент и начинаем заново.
   • Для сегмента длины $L$, если там действительно присутствуют и белые, и чёрные клетки, и $p \neq q$, число подотрезков длины $\ge 2$ равно $\frac{L(L-1)}{2}$.

2. Обработка нескольких строк

   • Если высота подматрицы $\ge 2$, то условие «есть белые и чёрные клетки» уже выполняется автоматически (поскольку в вертикальном столбце непременно встретятся и белая, и чёрная клетка).
   • Нужно лишь проверить, что все белые клетки внутри выбранных строк и столбцов — одной буквы $p$, а все чёрные — одной буквы $q$, причём $p \neq q$.
   • Фиксируем пару строк $(r_1, r_2)$, тогда высота $\ge 2$. Для каждого столбца проверяем «однородность» белых и чёрных клеток отдельно. Получаем пару $(p_c, q_c)$. Если $(p_c, q_c)$ совпадает подряд на отрезке столбцов, значит там все соответствующие клетки будут одинаковы, и весь блок столбцов даёт правильные подматрицы (для строк $[r_1..r_2]$). Количество способов выбрать подотрезок столбцов длины $L$ есть $\frac{L(L+1)}{2}$.

3. Сложность

   • Перебор всех пар строк $(r_1, r_2)$ даёт $O(n^2)$.
   • Для каждой пары строк мы за $O(m)$ колонок делаем проверку «однородности» для белых и чёрных клеток. Но сама проверка «однородности» в наивном виде требовала бы посмотреть каждую клетку. Чтобы не делать $O(m \times (r_2-r_1+1))$, мы пользуемся заранее вычисленными префиксными частотами для каждой буквы (и для белых/чёрных раздельно), что позволяет за $O(26)$ проверить, какая именно буква заполнена, либо понять, что их больше одной. Итого получается $O(n^2 \times m \times 26)$.
   • Хотя формально при $n,m=300$ это может быть весьма большим числом операций, на практике в скомпилированных языках (C++/Rust/Go) подобный код часто укладывается в разумные секунды. В Python же придётся либо очень оптимизировать, либо иметь большой лимит по времени.

4. Проверка на примерах

   • Пример 1:

     text
     2 2
     ab
     ba

     Ответ должен быть 5.

     • Однострочные подматрицы:
       • Строка 0: "ab" даёт один отрезок длины 2 ⇒ $\frac{2\cdot 1}{2}=1$.
       • Строка 1: "ba" аналогично 1. Итого 2.
     • Двухстрочные подматрицы (r1=0,r2=1):
       • Столбец 0 проверяем белые клетки: (0,0)->'a', (1,0)->'b' (но (1,0) чёрная!), значит по белым смотрим (0,0) и (1,1), однако (1,1) в другом столбце, так что чисто в столбце 0 белые — это (0,0): 'a'. По чёрным (1,0): 'b'. Аналогично столбец 1 по белым (1,1): 'a', по чёрным (0,1): 'b'. Оба столбца дают пару (p='a',q='b'). Идёт подряд из 2 столбцов ⇒ $\frac{2\cdot 3}{2}=3$. Суммарно $2+3=5$.

     Всё сходится с условием.

   • Пример 2:

     text
     2 2
     aa
     aa

     Здесь ответ 0, потому что хоть клетки и по цветам («белые» = 'a', «чёрные» = 'a'), но это одна и та же буква, а нужно ровно две разные. Проверка кодом покажет, что $p=q$ (оба равны 'a'), значит подматрицы не засчитываются.

Таким образом, данный код (при наличии достаточных ресурсных лимитов) решает задачу корректно. Если же время/память слишком ограничены для Python, обычно переходят на C++ и тщательно оптимизируют реализацию (используют быструю работу с массивами, убирают лишние проверки и т.п.).