CODEWITHSUNDEEP
c++stdstdvectorprefix-sum
Y.T.
Y.T.

Reputation: 2748

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.

Upvotes: 4

Views: 1614

Answers (1)

Caleth
Caleth

Reputation: 63039

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());
}

Upvotes: 2

Related Questions