User purchase platform

# Table: Spending
# 
# +-------------+---------+
# | Column Name | Type    |
# +-------------+---------+
# | user_id     | int     |
# | spend_date  | date    |
# | platform    | enum    |
# | amount      | int     |
# +-------------+---------+
# The table logs the history of the spending of users that make purchases from an online shopping website that has a desktop and a mobile application.
# (user_id, spend_date, platform) is the primary key (combination of columns with unique values) of this table.
# The platform column is an ENUM (category) type of ('desktop', 'mobile').
# 
#  
# 
# Write a solution to find the total number of users and the total amount spent using the mobile only, the desktop only, and both mobile and desktop together for each date.
# 
# Return the result table in any order.
# 
# The result format is in the following example.
#
# Example 1:
# Input:
# Spending table:
# +---------+------------+----------+--------+
# | user_id | spend_date | platform | amount |
# +---------+------------+----------+--------+
# | 1       | 2019-07-01 | mobile   | 100    |
# | 1       | 2019-07-01 | desktop  | 100    |
# | 2       | 2019-07-01 | mobile   | 100    |
# | 2       | 2019-07-02 | mobile   | 100    |
# | 3       | 2019-07-01 | desktop  | 100    |
# | 3       | 2019-07-02 | desktop  | 100    |
# +---------+------------+----------+--------+
# Output:
# +------------+----------+--------------+-------------+
# | spend_date | platform | total_amount | total_users |
# +------------+----------+--------------+-------------+
# | 2019-07-01 | desktop  | 100          | 1           |
# | 2019-07-01 | mobile   | 100          | 1           |
# | 2019-07-01 | both     | 200          | 1           |
# | 2019-07-02 | desktop  | 100          | 1           |
# | 2019-07-02 | mobile   | 100          | 1           |
# | 2019-07-02 | both     | 0            | 0           |
# +------------+----------+--------------+-------------+
# Explanation:
# On 2019-07-01, user 1 purchased using both desktop and mobile, user 2 purchased using mobile only and user 3 purchased using desktop only.
# On 2019-07-02, user 2 purchased using mobile only, user 3 purchased using desktop only and no one purchased using both platforms.

import pandas as pd

def user_purchase(spending: pd.DataFrame) -> pd.DataFrame:
  user_date = spending.groupby(['user_id', 'spend_date'])['platform'].apply(set).reset_index()
  user_date['platform_type'] = user_date['platform'].apply(
    lambda x: 'both' if len(x) == 2 else list(x)[0]
  )
  spending_with_type = spending.merge(
    user_date[['user_id', 'spend_date', 'platform_type']], on=['user_id', 'spend_date']
  )
  agg = spending_with_type.groupby(['spend_date', 'platform_type']).agg(
    total_amount=('amount', 'sum'),
    total_users=('user_id', 'nunique')
  ).reset_index()
  agg.columns = ['spend_date', 'platform', 'total_amount', 'total_users']
  dates = spending['spend_date'].unique()
  platforms = ['desktop', 'mobile', 'both']
  all_combos = pd.DataFrame([(d, p) for d in dates for p in platforms], columns=['spend_date', 'platform'])
  result = all_combos.merge(agg, on=['spend_date', 'platform'], how='left').fillna(0)
  result['total_amount'] = result['total_amount'].astype(int)
  result['total_users'] = result['total_users'].astype(int)
  return result