[每日一练]通过shift移动函数实现连续数据的需求

该题目来源于力扣：

603. 连续空余座位 - 力扣（LeetCode）

题目要求：

表: Cinema

+-------------+------+
| Column Name | Type |
+-------------+------+
| seat_id     | int  |
| free        | bool |
+-------------+------+
Seat_id 是该表的自动递增主键列。
在 PostgreSQL 中，free 存储为整数。请使用 ::boolean 将其转换为布尔格式。
该表的每一行表示第 i 个座位是否空闲。1 表示空闲，0 表示被占用。
 

查找电影院所有连续可用的座位。

返回按 seat_id 升序排序 的结果表。

测试用例的生成使得两个以上的座位连续可用。

结果表格式如下所示。

 

示例 1:

输入: 
Cinema 表:
+---------+------+
| seat_id | free |
+---------+------+
| 1       | 1    |
| 2       | 0    |
| 3       | 1    |
| 4       | 1    |
| 5       | 1    |
+---------+------+
输出: 
+---------+
| seat_id |
+---------+
| 3       |
| 4       |
| 5       |
+---------+

思路流程：

shift() 函数用于沿着指定轴将数据向前或向后移动

关于shift函数的语法：

DataFrame.shift(periods=1, freq=None, axis=0, fill_value=None)
'''
periods: 整数，表示要移动的偏移量。正数表示向下移动，负数表示向上移动。默认为 1。
freq: DateOffset、DateOffset-like、timedelta 或字符串，用于移动时间序列数据，通常在处理时间序列数据时使用。
axis: 整数，表示要移动的方向，0 表示沿着行的方向移动，1 表示沿着列的方向移动。
fill_value: 在移动过程中产生缺失值时要填充的值。
'''

首先，我们要先进行数据的清洗，使用shift函数对数据的排序有要求。我们可以利用sort_values()函数对原数据进行排序

import pandas as pd

def consecutive_available_seats(cinema: pd.DataFrame) -> pd.DataFrame:
    data=cinema.sort_values('seat_id')

然后我们需要解决两个需求：1.要求相邻的数据的free列都是“1"，2.”要求数据是相邻的。这个需求我们可以通过切片查询实现。

条件1，查询出的数据必须为1，可以表示为

data=data[(data['free']==1)]

条件2，要求查询出的相邻数据为为1，可以表示为

data=data[((data['free']==1) & (data.shift(1)['free']==1]))]

这里的shift(1)表示默认按行的方向向上移动一个单位。可以表示为我们查询的每个数据和它上面相邻的一个单位的数据都为1.但是这个代码不完整，因为我们还没有查询每个数据和它下面一个单位的数据的关系。所以我们在切片中的条件可以再复制一份，用“或”，将两者连接起来，在后者的条件中，shift函数中的参数应该更改为-1，表示该数据的后一位。这样就可以查询出相邻的两个数据都为1的情况了。

data=data[((data['free']==1) & (data.shift(1)['free']==1)) | ((data['free']==1) & (data.shift(-1)['free']==1))]

最后将代码联合起来即可

代码实现：

import pandas as pd

def consecutive_available_seats(cinema: pd.DataFrame) -> pd.DataFrame:
    data=cinema.sort_values('seat_id')
    data=data[((data['free']==1) & (data.shift(1)['free']==1)) | ((data['free']==1) & (data.shift(-1)['free']==1))]
    return data[['seat_id']]