如何在Python中用类从DataFrame获取数据并随机选择卡车

我有如下DataFrame，希望用类来定义数据结构。请问如何通过变量随机选择卡车？例如当选中Truck 1时，从DataFrame中返回对应数据：idnum = 1、trucktype = 'int'、arrival = 1、productlist = Product1与Product2的和=(3+2,5+0,8+9)=(5,5,17)。

# Dataframe:
import pandas as pd
dictionary = {
    'TruckName': [1, 2, 3], 
    'Type': ['int','out','out'], 
    'Arrival': [1, 2, 1,],
    'Process': [4,7,8],
    'Product1': [3,5,8],
    'Product2': [2,0,9],
    }
df = pd.DataFrame(dictionary)

# define truck:
class Truck:
    def __init__(self,**kwargs):
        if 'idnum' in kwargs:
            self.idnum=kwargs['idnum'] 
        if 'truckType' in kwargs:
            self.truckType=kwargs['truckType']
        if 'productList' in kwargs:
            self.productList=kwargs['productList']
        if 'arrival' in kwargs:
            self.arrival=kwargs['arrival']

1. 完善Truck类 #

调整Truck类的初始化逻辑，添加类方法实现从DataFrame行数据直接创建实例，同时自动计算你需要的productList：

import pandas as pd
import random

class Truck:
    def __init__(self, idnum, truck_type, arrival, product_list):
        self.idnum = idnum
        self.truckType = truck_type
        self.arrival = arrival
        self.productList = product_list
    
    @classmethod
    def from_df_row(cls, row, df):
        # 计算所有卡车的Product1与Product2的和，生成元组
        all_product_sums = tuple(df['Product1'] + df['Product2'])
        return cls(
            idnum=row['TruckName'],
            truck_type=row['Type'],
            arrival=row['Arrival'],
            product_list=all_product_sums
        )
    
    def __repr__(self):
        # 自定义实例打印格式，方便查看结果
        return f"Truck(idnum={self.idnum}, truckType='{self.truckType}', arrival={self.arrival}, productList={self.productList})"

2. 实现随机选择逻辑 #

编写函数完成随机选卡车并返回对应Truck实例的功能：

def random_pick_truck(df):
    # 随机选一个卡车ID
    random_truck_id = random.choice(df['TruckName'].tolist())
    # 筛选出对应行数据
    selected_row = df[df['TruckName'] == random_truck_id].iloc[0]
    # 创建并返回Truck实例
    return Truck.from_df_row(selected_row, df)

3. 测试验证 #

# 初始化DataFrame
dictionary = {
    'TruckName': [1, 2, 3], 
    'Type': ['int','out','out'], 
    'Arrival': [1, 2, 1,],
    'Process': [4,7,8],
    'Product1': [3,5,8],
    'Product2': [2,0,9],
}
df = pd.DataFrame(dictionary)

# 测试选中Truck 1的场景
truck1_row = df[df['TruckName'] == 1].iloc[0]
truck1 = Truck.from_df_row(truck1_row, df)
print(truck1)
# 输出：Truck(idnum=1, truckType='int', arrival=1, productList=(5, 5, 17))

# 测试随机选择
random_truck = random_pick_truck(df)
print(random_truck)
# 随机输出某辆卡车的信息，比如：Truck(idnum=3, truckType='out', arrival=1, productList=(5, 5, 17))

可选调整 #

如果你的实际需求是仅返回选中卡车自身的Product1+Product2值，而非所有卡车的总和，只需修改from_df_row方法中的计算逻辑：

@classmethod
def from_df_row(cls, row, df):
    # 仅计算当前卡车的Product1+Product2
    product_sum = (row['Product1'] + row['Product2'],)
    return cls(
        idnum=row['TruckName'],
        truck_type=row['Type'],
        arrival=row['Arrival'],
        product_list=product_sum
    )

内容的提问来源于stack exchange，提问作者Koala