Pandas刷题（LeetCode 601.体育馆的人流量）-CSDN博客

描述：

表：Stadium

+---------------+---------+
| Column Name   | Type    |
+---------------+---------+
| id            | int     |
| visit_date    | date    |
| people        | int     |
+---------------+---------+
visit_date 是该表中具有唯一值的列。
每日人流量信息被记录在这三列信息中：序号 (id)、日期 (visit_date)、 人流量 (people)
每天只有一行记录，日期随着 id 的增加而增加

编写解决方案找出每行的人数大于或等于 100 且 id 连续的三行或更多行记录。

返回按 visit_date 升序排列 的结果表。

查询结果格式如下所示。

示例 1:

输入：
Stadium 表:
+------+------------+-----------+
| id   | visit_date | people    |
+------+------------+-----------+
| 1    | 2017-01-01 | 10        |
| 2    | 2017-01-02 | 109       |
| 3    | 2017-01-03 | 150       |
| 4    | 2017-01-04 | 99        |
| 5    | 2017-01-05 | 145       |
| 6    | 2017-01-06 | 1455      |
| 7    | 2017-01-07 | 199       |
| 8    | 2017-01-09 | 188       |
+------+------------+-----------+
输出：
+------+------------+-----------+
| id   | visit_date | people    |
+------+------------+-----------+
| 5    | 2017-01-05 | 145       |
| 6    | 2017-01-06 | 1455      |
| 7    | 2017-01-07 | 199       |
| 8    | 2017-01-09 | 188       |
+------+------------+-----------+
解释：
id 为 5、6、7、8 的四行 id 连续，并且每行都有 >= 100 的人数记录。
请注意，即使第 7 行和第 8 行的 visit_date 不是连续的，输出也应当包含第 8 行，因为我们只需要考虑 id 连续的记录。
不输出 id 为 2 和 3 的行，因为至少需要三条 id 连续的记录。

数据准备：

data = [[1, '2017-01-01', 10], [2, '2017-01-02', 109], [3, '2017-01-03', 150], [4, '2017-01-04', 99], [5, '2017-01-05', 145], [6, '2017-01-06', 1455], [7, '2017-01-07', 199], [8, '2017-01-09', 188]]
stadium = pd.DataFrame(data, columns=['id', 'visit_date', 'people']).astype({'id':'Int64', 'visit_date':'datetime64[ns]', 'people':'Int64'})

分析：

SQL见LeetCode_sql_day16（601.体育馆的人流量）_leetcode 601. 体育馆的人流量-CSDN博客P

pandas同理

①先根据id排序

②筛选出符合规则的数据再进行排序

③将两次排序作差

④如果是连续的那么差值相等对差值进行分组计数
df['r4'] = df.groupby('r3')['r3'].transform('count')
⑤筛选出计数大于等于三的数据

⑥选择题目要求的字段

代码：

def human_traffic(stadium: pd.DataFrame) -> pd.DataFrame:
    df = stadium.sort_values('id')
    df['rank'] = df['id'].rank()
    df = df[df['people']>=100]
    df['r2'] = df['id'].rank()
    df['r3'] = df['rank']-df['r2']
    df['r4'] = df.groupby('r3')['r3'].transform('count')
    df = df[df['r4']>= 3]
    df = df[['id','visit_date','people']].sort_values('id',ascending=True)
    return df