Work backwards. Start by deciding the last apple you will pick, then the second to last, etc.
import heapq
def solve(apples, H, U):
"""Solve apple-picking problem.
apples must be a sequence of (H, W) pairs.
"""
if U == 0:
return sum(x[1] for x in apples if x[0] <= H)
apples = sorted(apples, reversed=True)
# also creates a copy, to not modify caller's data
picked_weight = 0
available_weights = [] # stored negative for heapq
offset = U - H % U
if offset == U: offset = 0
top = offset - U
while (apples or available_weights) and top <= H:
if available_weights:
picked_weight += -heapq.heappop(available_weights)
top += U
else:
top += U * max(1, (apples[-1][0] - top) // U)
while apples and apples[0][0] <= top:
heapq.heappush(available_weights, -apples.pop()[1])
return picked_weight
Simple test:
def test(expected, apples, H, U):
actual = solve(apples, H, U)
if expected != actual:
print "expected=%r actual=%r | H=%r U=%r apples=%r" % (
expected, actual, H, U, apples)
test(45, [(82, 30), (91, 10), (93, 5), (94, 15)], 100, 10)
test(22, [( 1, 1), ( 2, 1), (81, 10), (82, 10)], 100, 10)
test(20, [( 1, 10), ( 2, 10), (11, 1)], 20, 10)
test(20, [(81, 10), (82, 10), (90, 5)], 100, 10)
test(1, [(2**31 - 1, 1)], 2**31 - 1, 1)
Someone requested C++, so here it is. It's nearly identical code and logic to the above Python, except for one change: C++ stdlib's heap functions work with the max value instead of the min, so no need for the negation. (I kept this self-contained, but utilities such as a heap adapter and container inserter will make the code easier to use.)
#include <cstdio>
#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>
struct Apple {
int h, w;
friend bool operator<(Apple const& a, Apple const& b) {
return a.h < b.h or (a.h == b.h and a.w < b.w);
}
friend bool operator>(Apple const& a, Apple const& b) {
return b < a;
}
friend std::ostream& operator<<(std::ostream& s, Apple const& v) {
s << '(' << v.h << ", " << v.w << ')';
return s;
}
};
template<class T, class C, T C::*M>
struct sum {
T operator()(T const& cur, C const& next) { return cur + next.*M; }
};
int solve(std::vector<Apple> apples, int H, int U) {
if (U == 0) {
return std::accumulate(apples.begin(), apples.end(), 0, sum<int, Apple, &Apple::w>());
}
std::sort(apples.begin(), apples.end(), std::greater<Apple>());
int picked_weight = 0;
std::vector<int> available_weights; // NOT stored negative, unlike Python code
int offset = U - H % U;
if (offset == U) offset = 0;
int top = offset - U;
while ((apples.size() or available_weights.size()) and top <= H) {
if (available_weights.size()) {
picked_weight += available_weights.front();
std::pop_heap(available_weights.begin(), available_weights.end());
available_weights.pop_back();
top += U;
}
else {
top += U * std::max(1, (apples.back().h - top) / U);
}
while (apples.size() and apples.back().h <= top) {
available_weights.push_back(apples.back().w);
std::push_heap(available_weights.begin(), available_weights.end());
apples.pop_back();
}
}
return picked_weight;
}
C++ tests:
template<int N>
void test(int expected, Apple (&apples)[N], int H, int U) {
std::vector<Apple> v (apples, apples + N);
int actual = solve(v, H, U);
if (expected != actual) {
std::printf("expected=%d actual=%d | H=%d U=%d apples=[",
expected, actual, H, U);
std::vector<Apple>::const_iterator i = v.begin();
if (i != v.end()) {
std::cout << *i;
for (++i; i != v.end(); ++i) {
std::cout << ", " << *i;
}
}
std::cout << "]" << std::endl;
}
}
int main() {
{
Apple data[] = {{82, 30}, {91, 10}, {93, 5}, {94, 15}};
test(45, data, 100, 10);
}
{
Apple data[] = {{ 1, 1}, { 2, 1}, {81, 10}, {82, 10}};
test(22, data, 100, 10);
}
{
Apple data[] = {{ 1, 10}, { 2, 10}, {11, 1}};
test(20, data, 20, 10);
}
{
Apple data[] = {{81, 10}, {82, 10}, {90, 5}};
test(20, data, 100, 10);
}
{
int n = 2147483647; // 2**31 - 1
Apple data[] = {{n, 1}};
test(1, data, n, 1);
}
}