План курса / Все задачи / Алгоритмы и структуры данных

Накопление в 2D

В этом уроке ты:
  • Узнаешь, как применять технику накопительных массивов в двумерных массивах.
  • Освоишь построение массивов сумм для строк, столбцов и диагоналей.
  • Решишь оптимально задачу уровня собеседования из крупной технологической компании.
Задача с собеседования

Время перейти на новый уровень сложности! Эта задача может занять до 30 минут на реальном интервью, и многим кандидатам не удаётся найти решение вовремя.

Условие

Дан двумерный массив целых чисел board. Нужно найти максимальную сумму клеток, которые находятся под боем шахматного ферзя, включая клетку, на которой он стоит.

Пример 1

Ввод: board =
  [[3, 2, 5, 1],
   [4, 9, 1, 3],
   [2, 9, 2, 1]]
Вывод: 40
Объяснение: максимальная сумма 40 получается, если поставить ферзя на клетку [1, 1] (индексация с 0).

Пример 2

Ввод: board = [[1, 2, 3]]
Вывод: 6
В чем сложность задачи?

Наивное решение выглядит так:

  1. Поставить ферзя в каждую клетку поля.
  2. Для каждой позиции подсчитать сумму всех клеток, которые он бьёт.
using System;
using System.Collections.Generic;
using System.Linq;

public class Solution
{
    public static int MaxQueenSum(List<List<int>> board)
    {
        int n = board.Count;
        int m = board[0].Count;

        int maxSum = 0;
        for (int i = 0; i < n; ++i)
        {
            for (int j = 0; j < m; ++j)
            {
                int currSum = 0;

                // Прибавляем сумму всех элементов в строке i
                currSum += board[i].Sum();

                // Прибавляем сумму всех элементов в колонке j
                for (int k = 0; k < n; ++k)
                {
                    currSum += board[k][j];
                }

                // Прибавляем сумму всех элементов в главной диагонали
                for (int k = -Math.Min(i, j); k < Math.Min(n - i, m - j); ++k)
                {
                    currSum += board[i + k][j + k];
                }

                // Прибавляем сумму всех элементов в побочной диагонали
                for (int k = -Math.Min(i, m - j - 1); k < Math.Min(n - i, j + 1); ++k)
                {
                    currSum += board[i + k][j - k];
                }

                // Элемент board[i][j] просуммирован 4 раза в предыдущих циклах
                // Вычитаем лишние разы
                currSum -= 3 * board[i][j];

                // Обновляем ответ
                if (i == 0 && j == 0)
                {
                    maxSum = currSum;  // Инициализируем maxSum, если это первая итерация
                }
                maxSum = Math.Max(maxSum, currSum);
            }
        }
        return maxSum;
    }
}

Однако такой подход имеет сложность по времени O(n * m * (n + m)), что неэффективно для больших массивов.

Оптимальное решение

Чтобы оптимизировать решение по времени до O(n * m), используем суммы для строк, столбцов и диагоналей.

Суммы для строк и столбцов

  • Суммы строк: rowsSum[i] — сумма всех элементов строки i.
  • Суммы столбцов: colsSum[j] — сумма всех элементов столбца j.

Подсчет rowsSum и colsSum в решении выглядит так:

function maxQueenSum(board) {
    // rows_sum[i] будет содержать сумму всех элементов для i-ой строки
    const rowsSum = new Array(board.length).fill(0);
    // cols_sum[j] будет содержать сумму всех элементов для j-ой колонки
    const colsSum = new Array(board[0].length).fill(0);

    for (let i = 0; i < board.length; i++) {
        for (let j = 0; j < board[i].length; j++) {
            const val = board[i][j];
            rowsSum[i] += val;
            colsSum[j] += val;
            // ...
        }
    }
}

Суммы для главных диагоналей

  • Главная диагональ: Индекс для неё вычисляется как (n - 1 - i) + j, где n - число строк. Подсчет d1Sum в решении выглядит так:
public class Solution
{
    public static int MaxQueenSum(List<List<int>> board)
    {
        // ...
        List<int> d1Sum = new List<int>(new int[board.Count + board[0].Count - 1]);

        for (int i = 0; i < board.Count; ++i)
        {
            for (int j = 0; j < board[i].Count; ++j)
            {
                int val = board[i][j];
                d1Sum[(board.Count - 1 - i) + j] += val;
            }
        }
        // ...
        return 0;
    }
}

Суммы для побочных диагоналей

  • Побочная диагональ: индекс в массиве d2Sum вычисляется как i + j .

Подсчет d2Sum в решении выглядит так:

public class Solution
{
    public static int MaxQueenSum(List<List<int>> board)
    {
        // ...
        List<int> d2Sum = new List<int>(new int[board.Count + board[0].Count - 1]);

        // делаем предподсчет, чтобы быстро находить позицию
        for (int i = 0; i < board.Count; ++i)
        {
            for (int j = 0; j < board[i].Count; ++j)
            {
                int val = board[i][j];
                d2Sum[i + j] += val;
            }
        }
        // ...
    }
}

Подсчет текущей суммы для каждой позиции ферзя

После построения массивов rowsSum, colsSum, d1Sum и d2Sum можно легко вычислить сумму клеток под боем ферзя:

  1. Для каждой клетки (i, j) суммируем: current_sum = rowsSum[i] + colsSum[j] + d1_sum[(len(board) - 1 - i) + j] + d2_sum[i + j]
  2. Вычитаем значение текущей клетки трижды, так как оно было учтено в каждой из сумм.
  3. Сравниваем currentSum с maxSum и обновляем максимум, если нужно.

Итоговое решение будет таким:

public class Solution
{
    public static int MaxQueenSum(List<List<int>> board)
    {
        // rowsSum[i] - будет содержать сумму всех элементов для i-ой строки
        List<int> rowsSum = new List<int>(new int[board.Count]);
        // colsSum[i] - будет содержать сумму всех элементов для i-ой колонки
        List<int> colsSum = new List<int>(new int[board[0].Count]);
        // d1Sum, d2Sum - содержат сумму элементов на диагоналях
        List<int> d1Sum = new List<int>(new int[board.Count + board[0].Count - 1]);
        List<int> d2Sum = new List<int>(new int[board.Count + board[0].Count - 1]);

        // делаем предподсчет, чтобы быстро находить позицию
        for (int i = 0; i < board.Count; ++i)
        {
            for (int j = 0; j < board[i].Count; ++j)
            {
                int val = board[i][j];
                rowsSum[i] += val;
                colsSum[j] += val;
                d1Sum[(board.Count - 1 - i) + j] += val;
                d2Sum[i + j] += val;
            }
        }

        int maxSum = 0;
        for (int i = 0; i < board.Count; ++i)
        {
            for (int j = 0; j < board[i].Count; ++j)
            {
                // при подсчете суммы складываем сумму текущей строки, столбца,
                // двух диагоналей и не забываем вычитать общий серединный элемент
                int currSum = rowsSum[i] + colsSum[j];
                currSum += d1Sum[(board.Count - 1 - i) + j] + d2Sum[i + j];
                currSum -= 3 * board[i][j];

                // обновляем максимальную сумму
                if (i == 0 && j == 0)
                {
                    // сумма может быть меньше 0, поэтому
                    // важно инициализировать начальным значением
                    maxSum = currSum;
                }
                maxSum = Math.Max(maxSum, currSum);
            }
        }
        return maxSum;
    }
}
Оценка по времени и памяти:
  • Время: O(n*m), где n — количество строк, m — количество столбцов.
  • Память: O(n + m), так как храним суммы для строк, столбцов и диагоналей.
Что дальше?

Теперь, когда ты освоил применение массивов сумм как в одномерном, так и в двумерном формате, не стесняйся применять свои знания на практике!