Very generic argmax function in C++ wanted

Question

I'm a spoiled Python programmer who is used to calculating the argmax of a collection with respect to some function with

max(collection, key=function)

For example:

l = [1,43,10,17]
a = max(l, key=lambda x: -1 * abs(42 - x))

a then contains 43, the number closest to 42.

Is it possible to write a C++ function which takes any "iterable" and any function and returns the argmax like above? I guess this would involve template parameters, the auto keyword, and range-based iteration, but I was not able to piece it together.

Answer 1

This is a two-step process. Define a function key which should get mapped to the elements, i.e. which is applied before the operation which finds the maximum. Wrap things together in a lambda expression defining the comparison for finding the maximum.

auto key = [](int x){
    return -abs(42 - x);
};

std::max_element(l.begin(), l.end(), [key](int a, int b){
    return key(a) < key(b);
});

Here, we have to capture key which was defined outside the second lambda function. (We could also have defined it inside). You can also put this in one single lambda function. When the 42 should be parameterized from outside the lambda, capture this as a variable:

int x = 42;
std::max_element(l.begin(), l.end(), [x](int a, int b){
    return -abs(x - a) < -abs(x - b);
});

Note that std::max_element returns an iterator. To access the value / a reference to it, prepend it with *:

int x = 42;
auto nearest = std::min_element(l.begin(), l.end(), [x](int a, int b){
    return abs(x - a) < abs(x - b);
});
std::cout << "Nearest to " << x << ": " << *nearest << std::endl;

You can nicely wrap this in a generic find_nearest function:

template<typename Iter>
Iter find_nearest(Iter begin, Iter end,
                  const typename std::iterator_traits<Iter>::value_type & value)
{
    typedef typename std::iterator_traits<Iter>::value_type T;
    return std::min_element(begin, end, [&value](const T& a, const T& b){
        return abs(value - a) < abs(value - b);
    });
}

auto a = find_nearest(l.begin(), l.end(), 42);
std::cout << *a << std::endl;

Live demo find_nearest: http://ideone.com/g7dMYI

---

A higher-order function similar to the argmax function in your question might look like this:

template<typename Iter, typename Function>
Iter argmax(Iter begin, Iter end, Function f)
{
    typedef typename std::iterator_traits<Iter>::value_type T;
    return std::min_element(begin, end, [&f](const T& a, const T& b){
        return f(a) < f(b);
    });
}

You can invoke this with the following code, having exactly the lambda function from your question:

auto a = argmax(l.begin(), l.end(), [](int x) { return -1 * abs(42 - x); });
std::cout << *a << std::endl;

Live demo argmax: http://ideone.com/HxLMap

---

The only remaining difference now is that this argmax function uses an iterator-based interface, which corresponds to the design of the C++ standard algorithms (<algorithm>). It's always a good idea to adapt your own coding style to the tools you're using.

If you want a container-based interface which returns the value directly, Nawaz provided a nice solution which requires the decltype-feature to correctly specify the return type. I decided to keep my version this way, so people can see the both alternative interface designs.

Answer 2

Since @leemes solutions are too many. All are correct, except that none attempts to imitate the Python version in your example, Here is my attempt to imitate that:

Convenient generic argmax-function just like Python version:

template<typename Container, typename Fn>
auto max(Container const & c, Fn && key) -> decltype(*std::begin(c))
{  
    if ( std::begin(c) == std::end(c) ) 
       throw std::invalid_argument("empty container is not allowed.");

    typedef decltype(*std::begin(c)) V;
    auto cmp = [&](V a, V b){ return key(a) < key(b); };
    return *std::max_element(std::begin(c), std::end(c), cmp);
}

And use it as:

std::vector<int> l = {1,43,10,17};
auto a = max(l, [](int x) { return -1 * std::abs(42-x); };

int l[] = {1,43,10,17}; //works with array also!
auto a = max(l, [](int x) { return -1 * std::abs(42-x); };

Note: Unlike the other solution, this max() returns the element itself, not the iterator to the element!

Also note this solution would work for user-defined container also:

namespace test
{
     template<size_t N>
     struct intcollection
     {
         int _data[N];
         int const * begin() const { return _data; }
         int const * end() const { return _data + N; }
     };
}

test::intcollection<4> c{{1,43,10,17}};
auto r = max(c, [](int x) { return -1 * std::abs(42-x); });

See the live demo.