Метод префиксных сумм

Метод префиксных сумм

Вадим Володин

Привет! Метод префиксных сумм - достаточно простой, но порой очень нужный алгоритм, о котором полезно знать. Я расскажу про две задачи, которые можно решить с помощью этого метода. Код примеров будет на языке Python.

ПРОСТАЯ ЗАДАЧА

Рассмотрим алгоритмическую задачу. Есть базар, на котором в ряд стоит N палаток. Каждая из палаток приносит определённое количество прибыли. Далее поступает Q запросов вида: а сколько прибыли приносят палатки начиная с палатки с номером l, заканчивая палаткой с номером r?

Простое решение

Пусть у нас есть запрос, l - левая граница, r - правая граница. В простом решении мы можем пройти по всем палаткам от l до r и посчитать их сумму.

queries = [0] * Q # Массив запросов, где запрос - пара (l, r)
profit = [0] * N # Массив количества прибыли для палаток

# ... 
# input queries and profit 
# ...

for l, r in queries:
  sum = 0
  for tent in range(l, r):
    sum += profit[tent]
  print(sum)

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

Алгоритмическая сложность: O(Q * N)

Это много. Можно подумать, как сделать быстрее.

Решение методом префиксных сумм

Метод префиксных сумм заключается в том, чтобы заранее преподсчитать сумму прибыли от первой палатки до каждой из палаток. Тогда получится массив:

prefix_sum[i] = profit[0] + profit[1] + ... + profit[i]

В таком случае ответом на запрос (l, r) является:

prefix_sum[r] - prefix_sum[l - 1]

Код будет выглядеть так:

prefix = [0] * N
for i in range(N):
  prefix[i] = profit[i] + prefix[i - 1]

for l, r in queries:
  print(prefix[r] - prefix[l - 1])

Это и есть метод префиксных сумм.

Для практики можно сдать эту задачу на codeforces.

BOX FILTER

Рассмотрим теперь другую задачу, из компьютерного зрения, которую тоже можно решить методом префиксных сумм.

Если рассматривать чёрно-белое изображение, можно представить его как матрицу пикселей, где значение каждого пикселя может быть от нуля (чёрный цвет) до 255 (белый цвет).

Box filter - один из алгоритмов в компьютерном зрении, который позволяет размыть изображение. В этом алгоритме для каждого пикселя записывается значение, равное среднему среди него самого и всех соседей.

Значение пикселя, выделенного красным, будет равно (2 + 2 + 5 + 3 + 1 + 5 + 3 + 3 + 5) / 9 = 29 / 9, можно округлить до трёх. Так делаем с каждым пикселем. В данном случае размер фильтра - 3 * 3, но он может быть и другого размера.

На этой картинке вы можете видеть результат применения box filter-а:


Простое решение

Рассмотрим, как выглядит простое решение этой задачи. Нужно пройтись по каждому пикселю, сложить значения всех соседних и поделить на их количество. Пусть N - размер квадратной картинки, M - размер квадратного фильтра.

for row in range(N):
  for column in range(N):
    sum = 0
    for row_adding in range(- M / 2, M / 2 + 1):
      for column_adding in range(- M / 2, M / 2 + 1):
        sum += image[row + row_adding][column + column_adding]
    new_image[row][column] = 1 / 9 * sum

Алгоритмическая сложность: O(N * N * M * M)

Решение методом префиксных сумм

С помощью метода префиксных сумм можно улучшить алгоритмическую сложность до O(N * N). Будем использовать двумерные префиксные суммы. Для этого необходимо для каждого пикселя (i, j) поддерживать сумму значений всех пикселей с индексом строки не более i и с индексом столбца не более j в двумерном массиве prefix.

Предположим, мы хотим посчитать всех соседних пикселей для текущего (i, j). На картинке такие пиксели в зелёном квадрате:

Такая сумма будет равна сумме значений по всему полю до нижнего правого угла текущего квадрата, за исключением того, что находится слева и сверху от него. Но если мы от всего поля до нижнего правого угла вычтем то, что слева и то, что сверху, понадобится прибавить обратно то, что одновременно и сверху и слева:

Код будет выглядеть так:

# Преподсчёт
for i in range(N):
  for j in range(N):
    prefix[i][j] = image[i][j] 
                    + prefix[i - 1][j] 
                    + prefix[i][j - 1] 
                    - prefix[i - 1][j - 1]

# Вычисление нового изображения
for i in range(N):
  for j in range(N):
    i0 = i - M / 2
    i1 = i + M / 2
    j0 = j - M / 2
    j1 = j + M / 2
    new_image[i][j] = prefix[i1][j1] 
                      - prefix[i1][j0 - 1] 
                      - prefix[i0 - 1][j1] 
                      + prefix[i0 - 1][j0 - 1]

Решение, оптимизированное по памяти

Предыдущий алгоритм быстрый, но дополнительно использует O(N * N) памяти. Фильтр обычно меньше изображения, поэтому можно сократить использование дополнительной памяти до O(N * M).

Для этого вернёмся к одномерным префиксным суммам и будем для каждой строки отдельно хранить префиксные суммы в массиве prefix_row. Будем хранить только M последних строк, чтобы экономить память и для обращения к текущей строке i будем использовать индекс i % M. Так prefix_row[i % M][j] будет хранить сумму image[i][0] + image[i][1] + ... + image[i][j].

Нужно понимать, что в массиве prefix_row ровно M строк, которые переиспользуются циклически. Так, для фильтра 3 * 3 строка с индексом ноль сначала будет использоваться для нулевого ряда изображения. Затем она будет использоваться для первого ряда, а строка с индексом два - для нулевого ряда. Далее строка с индексом ноль будет использоваться для второго ряда, а с индексом два - для первого ряда, с индексом один - для нулевого ряда. На следующей итерации строка с индексом ноль - для третьего ряда, с индексом один - для второго, с индексом два - для первого, а нулевой ряд изображения уже нигде не будет храниться.

Также в массиве prefix_2d будем хранить такую же сумму, но для M последних строк. Тогда для фильтра 3 * 3 значение пикселя (i, j) в новом изображении будет равняться prefix_2d[j + 1] - prefix_2d[j - 2].

Код будет выглядеть так:

for i in range(N):
  for j in range(N):
    prefix_2d[j] -= prefix_row[i % M][j]
    prefix_row[i % M][j] = image[i][j] + prefix_row[i % M][j - 1]
    prefix_2d[j] += prefix_row[i % M][j]
    new_image[i - M / 2][j - M / 2] = prefix_2d[j] 
                                    - prefix_2d[j - M - 1]


На этом всё. Мы рассмотрели применение метода префиксных сумм в решении двух задач.

В следующем посте я расскажу об ещё одной задаче компьютерного зрения, более сложной, в которой тоже можно применить метод префиксных сумм.

Больше постов вы можете посмотреть в телеграм канале @polyprogrammist_dev


Report Page