Search code examples
c++stdstdvectorprefix-sum

What is the cleanest way to do a `std::partial_sum` with a `0` in front?


In C++, there is a function std::partial_sum to compute prefix sum. The code

#include <iostream>
#include <vector>
#include <iterator>
#include <numeric>

int main() {
    std::vector<int> a = {1, 2, 3, 4, 5};
    std::partial_sum(a.begin(), a.end(), a.begin());
    return 0;
}

will override a to 1 3 6 10 15, which is expected.

However, in most cases I want to use prefix sum, I want a 0 in front to indicate "empty sum" so that I can use a[2] - a[0] to query the sum of first two elements. (That allows me to use a trivial nested for loop to find sum of all subarrays). Is there a way to achieve it with the std::partial_sum function? I don't know if it will be possible since the output size will be input size + 1.

Note: I am not seeking for ways that alters the content or type of a beforehand.

If the size of a is a concern:

#include <iostream>
#include <vector>
#include <iterator>
#include <numeric>

int main() {
    std::vector<int> a = {1, 2, 3, 4, 5, -1};
    std::partial_sum(a.begin(), a.end() - 1, a.begin());
    return 0;
}

Something like this can also work for me.


Solution

  • Is there a way to achieve it with the std::partial_sum function?

    Just write a 0 to the output iterator before calling std::partial_sum. Care should be taken as the output is one larger than the input, and this won't work in-place, as it writes the first output before reading the first input.

    template<class InputIt, class OutputIt>
    constexpr OutputIt my_partial_sum(InputIt first, InputIt last, OutputIt d_first)
    {
        *d_first++ = typename std::iterator_traits<InputIt>::value_type{};
        return std::partial_sum(first, last, d_first);
    }
    

    If you want to be able to do it in place, you could adapt the possible implementation of std::partial_sum further

    template<class InputIt, class OutputIt>
    constexpr OutputIt partial_sum(InputIt first, InputIt last, OutputIt d_first)
    {
        using value_type = typename std::iterator_traits<InputIt>::value_type;
    
        if (first == last) {
            *d_first++ = value_type{};
            return d_first;
        }
    
        value_type sum{};
        value_type next = *first;
    
        *d_first++ = sum;
     
        while (++first != last) {
           next = *first;
           sum = std::move(sum) + next;
           *d_first++ = sum;
        }
        return d_first;
    }
    

    but I think the simpler thing would be to prepend a 0 to your container.

    template <typename Container>
    void my_partial_sum(Container& c) {
        c.emplace(c.begin());
        std::partial_sum(c.begin(), std::prev(c.end()), c.begin());
    }