Редакция для Сколько дней ждать роста цены


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Автор: montes332

1. Идея

Для каждого дня нужно найти первый день справа, где цена строго больше текущей.

Если делать это напрямую, то для каждого i можно искать j > i перебором. Но тогда в худшем случае получится O(n^2), что слишком медленно при n до 2 * 10^5.

Здесь подходит монотонный стек.

Будем идти по дням слева направо и хранить в стеке индексы тех дней, для которых ответ пока не найден. Как только текущая цена p[i] оказывается больше цены на вершине стека, это значит, что для этого дня мы наконец нашли ближайший следующий день с большей ценой.

2. Наблюдения

  1. Если день j лежит в стеке, значит для него мы пока не встретили справа цену больше p[j].

  2. Когда пришёл новый день i, нужно проверить:

    • если p[i] > p[j] для вершины стека j, то именно день i является ответом для j, потому что:
      • мы идём слева направо;
      • раньше дня i подходящего большего значения не было, иначе j уже удалили бы из стека.
  3. В стеке удобно хранить индексы дней, а не сами цены, потому что в ответе нужно вывести разность индексов: i - j.

  4. Из стека удаляются все дни с меньшей ценой, чем текущая:

    • пока стек не пуст
    • и p[st.back()] < p[i]
    • извлекаем индекс j и пишем ans[j] = i - j
  5. Если цены равны, удалять нельзя, потому что нужна именно строго большая цена. Значит условие в цикле должно быть именно p[st.back()] < p[i], а не <=.

3. Алгоритм

  1. Считать n и массив цен p.
  2. Создать массив ответов ans длины n, заполненный нулями.
  3. Создать пустой стек st, в котором будут лежать индексы дней.
  4. Пройти по всем дням i от 0 до n - 1:
    • пока стек не пуст и цена на вершине меньше текущей цены:
      • извлечь индекс j из стека;
      • записать ans[j] = i - j;
    • добавить индекс i в стек.
  5. После завершения прохода в стеке останутся дни, для которых справа не нашлось большей цены. Для них ответ уже равен 0, ничего делать не нужно.
  6. Вывести массив 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)

Комментарии

Еще нет ни одного комментария.