Prefix Sum
October 3, 2024·Johnathan Jacobs
Used to find the sum of elements in multiple windows of an array more efficiently. Instead of summing the window (and redoing work) for each variation of the window on the array, we create an array where each element represents the sum of all previous elements of the source array. We then use this array to calculate the sums for the windows instead.
flowchart LR; subgraph A["Array"] direction LR 1---2---3---4---5---6---7 end
The prefix sum (calculated with P[i] = A[0] + ... + A[i]
) would be:
flowchart LR; subgraph P["Prefix Array"] direction LR 1---3---6---10---15---21---28 end
Now any now the sum in a window can be calculated with sum(i,j) = P[j]-P[i-1]
This takes an $O(n*m)$ operation and makes it $O(n)$ at the cost of $O(n)$ space.
Implementation
Here’s a C++ implementation using nested loops to calculate the prefix sum:
#include <vector>
std::vector<int> windowSums(const std::vector<int>& data,
const std::vector<std::pair<std::size_t,std::size_t>> windows) {
const auto prefixSum = [&](){
const auto size = data.size();
std::vector<int> out(size);
for(std::size_t i = 0; i < size; ++i) {
for(std::size_t j = 0; j < i; ++j) {
out[i] += data[j];
}
}
return out;
}();
std::vector<int> sums;
sums.reserve(windows.size());
for(auto& window : windows) {
sums.push_back(prefixSum[window.second] - prefixSum[window.first - 1]);
}
return sums;
}
int main() {
return windowSums({1, 2, 3, 4, 5, 6, 7, 8},
{{1, 2}, {1, 3}, {1, 5}, {5, 7}}
)[0];
}
C++ provides std::partial_sum
that simplifies the implementation:
#include <vector>
#include <ranges>
#include <numeric> // for std::partial_sum
// NOTE: ideally std::partial_sum could be replaced with an std::ranges::inclusive_scan
std::vector<int> windowSums(const std::vector<int>& data,
const std::vector<std::pair<std::size_t, std::size_t>>& windows) {
// Calculate the prefix sums
std::vector<int> prefixSum(data.size() + 1, 0); // One extra element for easy subtraction
std::partial_sum(data.cbegin(), data.cend(), prefixSum.begin() + 1);
// Compute sums for each window in windows
return windows
| std::views::transform([&](auto&& window) {
return prefixSum[window.second] - prefixSum[window.first - 1];
})
| std::ranges::to<std::vector<int>>();
}
int main() {
return windowSums(
{1, 2, 3, 4, 5, 6, 7, 8},
{{1, 2}, {1, 3}, {1, 5}, {5, 7}}
)[1];
}