Search code examples
c++generatorstdc++23

Get n'th element yielded from a std::generator


I am trying to implement a function that computes n choose k. My idea is to generate the nth row of Pascal's triangle, and return its kth element. I would like to use a std::generator for this.

The problem is I cannot find a way to get just the nth yielded element.

My code looks something like this:

#include <generator>
#include <ranges>
#include <vector>

std::generator<std::vector<long>> pascal_triangle()
{
    std::vector<long> current_row{1};

    while (true)
    {
        co_yield current_row;

        std::vector<long> next_row;
        next_row.reserve(current_row.size() + 1);

        next_row.push_back(1);

        for (const auto& [left_element, right_element] : std::views::pairwise(current_row))
        {
            next_row.push_back(left_element + right_element);
        }

        next_row.push_back(1);

        current_row = std::move(next_row);
    }
}

long choose(const int n, const int k)
{
    // This line does not compile.
    const std::vector row = std::ranges::next(pascal_triangle(), n);

    return row[k];
}

What do I need to do to get the nth row?


Solution

  • There are two problems with this line:

    const std::vector row = ranges::next(pascal_triangle(), n);
    

    The first is that std::ranges::next takes an iterator, not a range. The second is that it returns an iterator, not an element. std::ranges::next(i, n) is a generalization of i + n, it's not a generalization of i[n].

    Fixing both issues, this compiles:

    const std::vector row = *std::ranges::next(pascal_triangle().begin(), n);
    

    Of course, you could write your own function that does what you had wanted ranges::next to do:

    template <ranges::input_range R>
    auto nth(R&& r, ptrdiff_t n) -> decltype(auto) {
       return *ranges::next(ranges::begin(r), n);
    }