Редакция для Сколько дней ждать роста цены
Submitting an official solution before solving the problem yourself is a bannable offence.
Автор:
1. Идея
Для каждого дня нужно найти первый день справа, где цена строго больше текущей.
Если делать это напрямую, то для каждого i можно искать j > i перебором. Но тогда в худшем случае получится O(n^2), что слишком медленно при n до 2 * 10^5.
Здесь подходит монотонный стек.
Будем идти по дням слева направо и хранить в стеке индексы тех дней, для которых ответ пока не найден. Как только текущая цена p[i] оказывается больше цены на вершине стека, это значит, что для этого дня мы наконец нашли ближайший следующий день с большей ценой.
2. Наблюдения
Если день
jлежит в стеке, значит для него мы пока не встретили справа цену большеp[j].Когда пришёл новый день
i, нужно проверить:- если
p[i] > p[j]для вершины стекаj, то именно деньiявляется ответом дляj, потому что:- мы идём слева направо;
- раньше дня
iподходящего большего значения не было, иначеjуже удалили бы из стека.
- если
В стеке удобно хранить индексы дней, а не сами цены, потому что в ответе нужно вывести разность индексов:
i - j.Из стека удаляются все дни с меньшей ценой, чем текущая:
- пока стек не пуст
- и
p[st.back()] < p[i] - извлекаем индекс
jи пишемans[j] = i - j
Если цены равны, удалять нельзя, потому что нужна именно строго большая цена. Значит условие в цикле должно быть именно
p[st.back()] < p[i], а не<=.
3. Алгоритм
- Считать
nи массив ценp. - Создать массив ответов
ansдлиныn, заполненный нулями. - Создать пустой стек
st, в котором будут лежать индексы дней. - Пройти по всем дням
iот0доn - 1:- пока стек не пуст и цена на вершине меньше текущей цены:
- извлечь индекс
jиз стека; - записать
ans[j] = i - j;
- извлечь индекс
- добавить индекс
iв стек.
- пока стек не пуст и цена на вершине меньше текущей цены:
- После завершения прохода в стеке останутся дни, для которых справа не нашлось большей цены. Для них ответ уже равен
0, ничего делать не нужно. - Вывести массив
ans.
4. Почему это работает
Докажем корректность.
Рассмотрим некоторый день j.
Когда ответ для j записывается
Ответ для j записывается в момент, когда мы рассматриваем некоторый день i, такой что p[i] > p[j], и индекс j находится на вершине стека или под элементами, которые уже будут удалены раньше.
В этот момент выполняется ans[j] = i - j.
Нужно понять, почему именно день i является ближайшим следующим днём с большей ценой.
Почему найденный день действительно подходит
Во время обработки дня i условие p[j] < p[i] выполняется, значит цена в день i строго больше цены в день j. Также i > j, потому что мы идём слева направо.
Следовательно, день i действительно подходит как кандидат на ответ для j.
Почему не существует более раннего подходящего дня
Предположим, что существует день k, где j < k < i и p[k] > p[j].
Тогда при обработке дня k индекс j уже должен был быть удалён из стека, и ответом для него стало бы k - j. Но этого не произошло: j дожил в стеке до момента обработки i.
Противоречие.
Значит, среди всех дней справа первый день с большей ценой — именно i.
Почему дни, оставшиеся в стеке, получают ответ 0
Если индекс j так и остался в стеке после обработки всех дней, значит справа от него не встретилось ни одного дня i, для которого p[i] > p[j].
Следовательно, для такого дня более высокой цены в будущем нет, и правильный ответ равен 0.
Итак, алгоритм правильно находит ответ для каждого дня.
5. Сложность
Каждый индекс:
- ровно один раз добавляется в стек;
- не более одного раза удаляется из стека.
Значит, суммарное число операций со стеком линейно.
Итоговая сложность:
- по времени:
O(n) - по памяти:
O(n)
6. Код на C++17
#include <iostream>
#include <vector>
using namespace std;
int main() {
int n;
cin >> n;
vector<long long> p(n);
for (int i = 0; i < n; i++) {
cin >> p[i];
}
vector<int> ans(n, 0);
vector<int> st;
for (int i = 0; i < n; i++) {
while (!st.empty() && p[st.back()] < p[i]) {
int j = st.back();
st.pop_back();
ans[j] = i - j;
}
st.push_back(i);
}
for (int i = 0; i < n; i++) {
if (i > 0) {
cout << ' ';
}
cout << ans[i];
}
cout << '\n';
return 0;
}
7. Код на Python 3
n = int(input())
p = list(map(int, input().split()))
ans = [0] * n
st = []
for i in range(n):
while st and p[st[-1]] < p[i]:
j = st.pop()
ans[j] = i - j
st.append(i)
print(*ans)
Комментарии