Як стати королем трейдингу на LeetCode. Розв’язуємо Best Time to Buy and Sell Stock

💡 Усі статті, обговорення, новини про Python — в одному місці. Приєднуйтесь до Python спільноти!

Всім привіт! Мене звати Валентин, я Software Engineer і у цій статті ми розглянемо серію популярних задач на LeetCode, пов’язаних із купівлею та продажем акцій.

З мого досвіду вони доволі часто зустрічаються на співбесідах. Розглянемо приклади задач з LeetCode, там представлено чотири задачі різних рівнів із цієї серії. Для написання коду будемо використовувати Python.

Поїхали!

Для початку пропоную розглянути другу задачу по рівню складності, потім буде зрозуміло чому. 🙂

Best Time to Buy and Sell Stock II

Умови:

Дано цілочисельний масив prices, де prices[i] — це ціна акції в i-й день.
Щодня Ви можете вирішити купити та/або продати акцію. Водночас Ви можете володіти не більше ніж однією акцією. Однак можна купити й одразу продати її в той самий день.

Знайдіть і поверніть максимальний прибуток, який можна отримати.

Приклад 1:

Вхід: prices = [7,1,5,3,6,4]
Вихід: 7
Пояснення:

• Купуємо на 2-й день (ціна = 1) і продаємо на 3-й день (ціна = 5), прибуток = 5 — 1 = 4.
• Потім купуємо на 4-й день (ціна = 3) і продаємо на 5-й день (ціна = 6), прибуток = 6 — 3 = 3.
• Загальний прибуток: 4 + 3 = 7.

Приклад 2:

Вхід: prices = [1,2,3,4,5]
Вихід: 4
Пояснення:

• Купуємо на 1-й день (ціна = 1) і продаємо на 5-й день (ціна = 5), прибуток = 5 — 1 = 4.
• Загальний прибуток: 4.

Приклад 3:

Вхід: prices = [7,6,4,3,1]
Вихід: 0
Пояснення:

• Немає можливості отримати позитивний прибуток, тому акцію не купуємо і максимальний прибуток = 0.

Посилання: leetcode.com/...​ell-stock-ii/description

Рішення:

Задача достатньо проста у вирішенні.

Давайте спочатку оцінимо, чи можливо купити та продати акцію в один і той самий день і отримати прибуток. В реальному житті — так, але у даній задачі — ні. Це, тому що ціна не змінюється упродовж дня. Тому одразу відкидаємо даний варіант.

Залишається лише варіант — купити дешевше в один день і продати дорожче в інший. Пропоную подивитись більш детально на приклад 1 і 2.

У першому прикладі ми купуємо по низькій ціні та продаємо на наступний день по високій. А в другому випадку купуємо на 1-й день і продаємо аж в останній 5-й.

Другий приклад ми можемо інтерпретувати за схожим алгоритмом як і в першому прикладі. Купувати сьогодні дешевше і продавати завтра дорожче. У такому випадку, замість однієї транзакції будемо мати 5 транзакцій. Але результат буде той самий — 4.

Виходить, що нам потрібно проітеруватись всім списком і просумувати різниці між сусідніми елементами. Але, різниця елементів повинна бути лише додатною. Тобто, ми сумуємо лише додатні різниці, всі інші просто ігноруємо. Звучить не складно, напишемо код:

class Solution:
   def maxProfit(self, prices: List[int]) -> int:
       result = 0
       previous_price = prices[0]
       for price in prices[1:]:
           if price > previous_price:
               result += price - previous_price
           previous_price = price
       return result

Складність алгоритму у нас вийшла лінійна — O(n). Немає додаткового використання памʼяті, окрім декількох додаткових змінних, але ми їх не будемо рахувати, бо на великій кількості елементів, вони ні на що не впливають.

Розглянемо наступні три задачі, вони є дуже схожими. Почнемо з найпростішої.

Best Time to Buy and Sell Stock

Умови:

Вам дано масив prices, де prices[i] — це ціна акції в i-й день.

Ви хочете максимізувати свій прибуток, вибравши один день для покупки акції та інший день у майбутньому для її продажу.

Поверніть максимальний прибуток, який можна отримати від цієї угоди. Якщо прибуток отримати неможливо, поверніть 0.

Приклад 1:

Вхід: prices = [7,1,5,3,6,4]
Вихід: 5
Пояснення:

• Купуємо на 2-й день (ціна = 1) і продаємо на 5-й день (ціна = 6), прибуток = 6 — 1 = 5.
• Зверніть увагу, що купівля на 2-й день і продаж на 1-й день неможливі, оскільки спочатку потрібно купити, а потім продавати.

Приклад 2:

Вхід: prices = [7,6,4,3,1]
Вихід: 0
Пояснення:

• У цьому випадку неможливо отримати позитивний прибуток, тому жодних угод не здійснюється, і максимальний прибуток дорівнює 0.

Посилання: leetcode.com/...​d-sell-stock/description

Рішення:

З першого погляду, задача звучить простіше, потрібно лише зробити 1 транзакцію з максимальним профітом.

Одразу приходить ідея знайти мінімум і максимум у списку і просто порахувати різницю. Але, мінімум повинен знаходитись до максимум у списку, бо ми не можемо спочатку продати, а потім купити акцію. Що тоді в цьому випадку?

Шукати наступний максимум, який менше попереднього? А якщо навпаки, потрібно шукати наступний мінімум? Цей варіант нам не підходить, потрібно обрати інше рішення.

Наступний варіант, який приходить в голову, це взяти перший елемент списку і врахувати його як мінімум і пошукати у решті списку — максимальний елемент. Потім порахувати різницю, далі взяти наступний елемент списку, тільки якщо він менше вже існуючого мінімуму, і пошукати знову максимальний елемент у решті списку. Знову порахувати різницю і порівняти із вже існуючою різницею, і зберегти більшу.

Таким чином можна повторювати до поки ми не дійдемо до перед останнього елементу.

Спробуймо реалізувати даний алгоритм:

import math
class Solution(object):
   def maxProfit(self, prices: List[int]) -> int:
       min_value = math.inf
       diff = 0
       for ind, val in enumerate(prices[:-1]):
           if val < min_value:
               min_value = val
               new_diff = max(prices[ind + 1:]) - min_value
               if new_diff > diff:
                   diff = new_diff
       return diff

Наче алгоритм працюючий, тести проходять. Але давайте подивимось на складність даного алгоритму. В середині нашого головного циклу є ще вкладений цикл, який шукає максимум у меншому списку. А це по своїй суті є показником квадратичної залежності — O(n²). Щодо використання памʼяті, виглядає добре, але зі складністю потрібно щось робити.

Ця задача дуже нагадує задачу пошуку підсписку з найбільшою сумою елементів. Вирішується така задача за допомогою Kadane’s Algorithm.

Детальніше про сам алгоритм можна почитати тут: medium.com/...​does-it-work-3fd8849ed73d

Спробуймо адаптувати його до нашої задачі.

Спочатку потрібно знайти мінімум, а потім вже максимум і різницю. Якщо ми будемо ітеруватись нашим списком, то локальний мінімум (мінімум у нашому підсписку від 0 до і-го елементу) можна знайти простим порівнянням 2-х чисел, і-го елементу і мінімум і-1 підсписку.

Якщо коротко, ми спочатку вважаємо, що перший елемент є мінімальним. Далі порівнюємо 2-й елемент із підсписком із першого елементу. Якщо 2-й менший, то він стає мінімальним, якщо ні, то перший лишається мінімальним. І так далі.

З мінімумом розібрались. Але як знайти найбільший профіт для 1-ї транзакції? Все просто. Коли ми будемо ітеруватись списком, потрібно просто від кожного елементу віднімати наш вже існуючий локальний мінімум.

Таким чином, ми завжди будемо впевнені, що менше число знаходиться зліва і воно є мінімальним для підсписку зліва. Це саме те що нам треба. Нумо імплементувати!

import math
class Solution(object):
   def maxProfit(self, prices: List[int]) -> int:
       buy_price = math.inf
       diff = 0
       for price in prices:
           buy_price = min(buy_price, price)
           diff = max(diff, price - buy_price)
       return diff

Виглядає дуже просто. Якщо оцінити складність, то вона тут лінійна — O(n). Додаткової памʼяті ми не використовуємо.

Супер, це те що нам і потрібно!

Тепер перейдемо до найскладніших задач. Вони оцінюються як hard на LeetCode.

Best Time to Buy and Sell Stock III

Умови:

Вам дано масив prices, де prices[i] — це ціна акції в i-й день. Знайдіть максимальний прибуток, який можна отримати. Ви можете виконати не більше двох транзакцій.

Примітка: Ви не можете здійснювати кілька транзакцій одночасно (тобто спочатку потрібно продати акцію, перш ніж купити її знову).

Приклад 1:

Вхід: prices = [3,3,5,0,0,3,1,4]
Вихід: 6
Пояснення:

• Купуємо на 4-й день (ціна = 0) і продаємо на 6-й день (ціна = 3), прибуток = 3 — 0 = 3.
• Потім купуємо на 7-й день (ціна = 1) і продаємо на 8-й день (ціна = 4), прибуток = 4 — 1 = 3.
• Загальний прибуток: 3 + 3 = 6.

Приклад 2:

Вхід: prices = [1,2,3,4,5]
Вихід: 4
Пояснення:

• Купуємо на 1-й день (ціна = 1) і продаємо на 5-й день (ціна = 5), прибуток = 5 — 1 = 4.
• Зверніть увагу, що не можна купити на 1-й день, потім знову купити на 2-й день і продати пізніше, оскільки це означає одночасне виконання кількох транзакцій. Спочатку потрібно продати перед наступною покупкою.

Приклад 3:

Вхід: prices = [7,6,4,3,1]
Вихід: 0
Пояснення:

• У цьому випадку жодної вигідної транзакції здійснити не можна, тому максимальний прибуток = 0.

Посилання: leetcode.com/...​ll-stock-iii/description

Рішення:

Задача практично ідентична з попередньою, тільки тепер ми повинні робити не одну транзакцію, а одну чи дві.

Теоретично, ми можемо використати попередній алгоритм двічі та таким чином знайдемо 2 найбільш профітні транзакції. Тільки після першого разу потрібно видалити підсписок з нашого списку. Мінус даного підходу в тому, що нам потрібне додаткове використання памʼяті, бо необхідно буде створювати додатковий список.

По складності це буде лінійна залежність, але це тільки для даного випадку, де ми шукаємо дві транзакції. Спробуймо обійтись без додаткового використання памʼяті.

Подивимось детальніше на реалізацію Kadane’s Algorithm в попередній задачі та ще раз передивимось умови задачі. Нам не потрібно знайти саме дві транзакції з найбільшим профітом. Нам підійде або одна транзакція або дві, але вони повинні мати в сумі найбільший профіт.

Отже, нам потрібно врахувати профіт найбільшої транзакції та намагатись знайти наступну транзакцію із вже меншим профітом. Звучить взагалі дуже просто.

Беремо вираз з мінімальним значенням і просто враховуємо різницю попередньої (найбільшої) транзакції. Таким чином, рухаючись по списку ми завжди будемо мати профіт з найбільшої транзакції плюс профіт з наступною меншою транзакцією, якщо така існує. Ближче до коду, буде більш зрозуміло.

class Solution:
   def maxProfit(self, prices: List[int]) -> int:
       min_price_first_transaction = prices[0]
       min_price_second_transaction = prices[0]
       diff_first_transaction = 0
       diff_second_transaction = 0
       for ind in range(1, len(prices)):
           price = prices[ind]
           min_price_first_transaction = min(min_price_first_transaction, price)
           diff_first_transaction = max(
               diff_first_transaction,
               price - min_price_first_transaction
           )
          
           min_price_after_second_transaction = min(
               min_price_after_second_transaction,
               price - diff_first_transaction
           )
           diff_after_second_transaction = max(
               diff_after_second_transaction,
               price - min_price_after_second_transaction
           )
       return diff_after_second_transaction

Різниці майже немає, ті самі вирази для другої транзакції, але ми враховуємо профіт першої транзакції (diff_first_transaction) в мінімумі другої.

Таким чином ми або знаходимо максимальну різницю 2-х найбільших транзакцій, або профіт лише одної транзакції. І це все буде накопичуватись у змінній diff_after_second_transaction.

Тепер ми маємо найкращу складність алгоритму — O(n). І ми не використовуємо додаткової памʼяті.

Ну що, перейдемо до останньої задачі із цієї серії.

Best Time to Buy and Sell Stock IV

Умови:

Вам дано цілочисельний масив prices, де prices[i] — це ціна акції в i-й день, і ціле число k. Знайдіть максимальний прибуток, який можна отримати. Ви можете виконати не більше k транзакцій: тобто купити не більше k разів і продати не більше k разів.

Примітка: Ви не можете здійснювати кілька транзакцій одночасно (тобто спочатку потрібно продати акцію, перш ніж купити її знову).

Приклад 1:

Вхід: k = 2, prices = [2,4,1]
Вихід: 2
Пояснення:

• Купуємо на 1-й день (ціна = 2) і продаємо на 2-й день (ціна = 4), прибуток = 4 — 2 = 2.

Приклад 2:

Вхід: k = 2, prices = [3,2,6,5,0,3]
Вихід: 7
Пояснення:

• Купуємо на 2-й день (ціна = 2) і продаємо на 3-й день (ціна = 6), прибуток = 6 — 2 = 4.
• Потім купуємо на 5-й день (ціна = 0) і продаємо на 6-й день (ціна = 3), прибуток = 3 — 0 = 3.
• Загальний прибуток: 4 + 3 = 7.

Посилання: leetcode.com/...​ell-stock-iv/description

Рішення:

Це така сама задача як і попередня, але ми повинні просто параметризувати кількість транзакцій.

Найпростішим варіантом зробити це буде створити масив із мінімальних цін та профітів залежно від параметра k, який надісланий у функцію.

Переходимо до коду.

class Solution:
   def maxProfit(self, k: int, prices: List[int]) -> int:
       min_prices = [prices[0]] * k
       diff_transactions = [0] * k
      
       for ind in range(1, len(prices)):
           price = prices[ind]
           for transaction_ind in range(k):
               min_price = price 
               if transaction_ind > 0:
                   min_price = price - diff_transactions[transaction_ind-1]
              
               min_prices[transaction_ind] = min(
                   min_prices[transaction_ind], min_price
               )
               diff_transactions[transaction_ind] = max(
                   diff_transactions[transaction_ind], price - min_prices[transaction_ind]
               )
       return diff_transactions[-1]

Складність даного алгоритму — O(kn). Додаткова памʼять — це два списки довжиною k. З цікавого — ми можемо знайти профіт на кожну транзакцію, якщо проітеруємось по списку diff_transactions і відніматимемо попереднє значення. Комбінуючи це зі списком із мінімальних значень — ми зможемо знайти ціни покупок і продажу акцій на кожну транзакцію.

Висновок:

Тепер повернемось до питання яке мали на початку. Чому ж все ж таки ми почали розгляд задач саме з другої, а не першої? Бо саме перша, третя та четверта задачі вирішуються одним і тим самим алгоритмом із додатковим тюнингом. Тому ці задачі були згруповані разом.

Третя та четверта задачі оцінюються як hard рівень на LeetCode. І Ви бачите, як можна не дуже складно, розширивши алгоритм для простої задачі, вирішувати складніші. Сподіваюсь було цікаво і пізнавально! 🙂