IOI 1991 - Problem 2: Mangoes

// IOI 1991 - Problem 2: Mangoes
// Weighted job scheduling: each tree has interval [d, d+k-1] with value a.
// Pick at most one tree per day to maximize total harvest.
// DP with binary search, sorted by deadline. O(n log n)
#include <bits/stdc++.h>
using namespace std;

int main() {
    int n, k;
    scanf("%d%d", &n, &k);

    struct Tree {
        int start, end, value;
    };
    vector<Tree> trees(n);
    for (int i = 0; i < n; i++) {
        int d, a;
        scanf("%d%d", &d, &a);
        trees[i] = {d, d + k - 1, a};
    }

    // Sort by end time (deadline)
    sort(trees.begin(), trees.end(),
         [](const Tree& a, const Tree& b) { return a.end < b.end; });

    // dp[i] = max mangoes considering first i trees
    vector<long long> dp(n + 1, 0);

    for (int i = 1; i <= n; i++) {
        dp[i] = dp[i - 1]; // skip tree i-1

        // Binary search: find last tree j whose end < tree i's start
        int lo = 0, hi = i - 1, best = 0;
        while (lo <= hi) {
            int mid = (lo + hi) / 2;
            if (trees[mid].end < trees[i - 1].start) {
                best = mid + 1;
                lo = mid + 1;
            } else {
                hi = mid - 1;
            }
        }
        dp[i] = max(dp[i], dp[best] + trees[i - 1].value);
    }

    printf("%lld\n", dp[n]);
    return 0;
}

\documentclass[11pt,a4paper]{article}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb,amsthm}
\usepackage{listings}
\usepackage[margin=2.5cm]{geometry}
\usepackage{xcolor}

\newtheorem{theorem}{Theorem}
\newtheorem{lemma}[theorem]{Lemma}

\lstset{
  language=C++,
  basicstyle=\ttfamily\small,
  keywordstyle=\color{blue}\bfseries,
  commentstyle=\color{green!50!black},
  stringstyle=\color{red!70!black},
  numbers=left,
  numberstyle=\tiny\color{gray},
  breaklines=true,
  frame=single,
  tabsize=4,
  showstringspaces=false
}

\title{IOI 1991 -- Problem 2: Mangoes}
\author{Solution}
\date{}

\begin{document}
\maketitle

\section{Problem Statement}

There are $n$ mango trees. Tree $i$ has $a_i$ mangoes that ripen on day
$d_i$ and fall off after $k$~days. That is, tree $i$'s mangoes can be
harvested on any single day in the interval $[d_i,\, d_i + k - 1]$.

A harvester picks one tree per day, collecting all of its available mangoes.
Maximize the total number of mangoes harvested.

\section{Solution}

This is a \textbf{weighted interval scheduling} problem. Each tree defines
a job with interval $[d_i,\, d_i + k - 1]$ and weight $a_i$. We must
select a set of jobs (at most one per day) to maximize total weight.

\subsection{Greedy with Max-Heap}

Sweep through days from the earliest ripening day to the latest deadline:

\begin{enumerate}
  \item Maintain a max-heap of available trees, keyed by mango count.
  \item On each day, add all trees that ripen on or before this day.
  \item Remove expired trees (deadline passed) from the top of the heap.
  \item Pick the tree with the most mangoes and add its value to the total.
\end{enumerate}

\begin{lemma}[Correctness]
The greedy strategy of always picking the highest-value available job is
optimal for unit-time jobs on a timeline. This follows from the exchange
argument: if an optimal solution picks a lower-value job on some day while
a higher-value job is available, swapping them does not decrease the total.
\end{lemma}

\textbf{Caveat.} The greedy approach as stated has a subtle issue: popping
expired entries only from the top of the heap may leave expired entries
deeper in the heap. These stale entries are harmless because they will be
discarded when they eventually reach the top. However, if the heap contains
only expired entries on some day, the algorithm correctly produces no harvest
for that day.

\subsection{Potential Issue with Day-by-Day Sweep}

If $d_{\max} + k$ is very large (up to $10^9$), iterating over every day is
too slow. In that case, we should use event-driven processing: sort all trees
by ripening day, and for each day that has at least one event (a tree
ripening or an opportunity to harvest), process the heap. The number of
meaningful events is at most $n$, so the total work is $O(n \log n)$.

\section{Complexity Analysis}

\begin{itemize}
  \item \textbf{Time:} $O(n \log n)$ for sorting and heap operations.
  \item \textbf{Space:} $O(n)$ for the heap.
\end{itemize}

\section{Implementation}

The following implementation uses an event-driven approach to avoid iterating
over empty days.

\begin{lstlisting}
#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
using namespace std;

int main() {
    int n, k;
    cin >> n >> k;

    vector<pair<int,int>> trees(n); // (ripen_day, mangoes)
    for (int i = 0; i < n; i++)
        cin >> trees[i].first >> trees[i].second;

    // Sort by ripen day
    sort(trees.begin(), trees.end());

    // Max-heap of (mangoes, deadline)
    priority_queue<pair<int,int>> pq;
    long long total = 0;
    int idx = 0;

    // Process day by day, but skip empty days.
    // The meaningful days are: each tree's ripen day through its deadline.
    // We process trees greedily, one per day.
    int day = trees[0].first;
    int maxDay = 0;
    for (auto& t : trees)
        maxDay = max(maxDay, t.first + k - 1);

    while (day <= maxDay) {
        // Add all trees ripening on or before this day
        while (idx < n && trees[idx].first <= day) {
            pq.push({trees[idx].second, trees[idx].first + k - 1});
            idx++;
        }

        // Remove expired trees from top
        while (!pq.empty() && pq.top().second < day)
            pq.pop();

        // Pick the best available tree
        if (!pq.empty()) {
            total += pq.top().first;
            pq.pop();
        }

        // Skip to next meaningful day
        if (pq.empty() && idx < n)
            day = trees[idx].first;
        else
            day++;
    }

    cout << total << endl;
    return 0;
}
\end{lstlisting}

\section{Example}

Suppose $n = 3$, $k = 2$, with trees: $(d=1, a=5)$, $(d=2, a=3)$,
$(d=2, a=8)$.

\begin{itemize}
  \item Day 1: Tree 1 available. Pick tree 1 (5 mangoes).
  \item Day 2: Trees 2 and 3 available. Pick tree 3 (8 mangoes).
  \item Day 3: Tree 2 still available (deadline is day 3). Pick tree 2
    (3 mangoes).
\end{itemize}

Total: $5 + 8 + 3 = 16$.

\end{document}