在python数据帧中，当列a等于某个值时，如何找到列b

Question

我有一个包含用户id、电影id和评分的csv，我想返回某些用户id的所有电影id。这是数据的样子

userId  movieId  rating   timestamp
1        1     4.0   964982703
1        3     4.0   964981247
1        6     4.0   964982224
1       47     5.0   964983815
1       50     5.0   964982931

假设我想返回用户id为1的所有电影，并将它们存储在某种结构中，可能是一个列表。输出为[1,3,6,47,50，…]

我尝试使用这段代码只是为了打印输出

for x in ratings["userId"]:
for y in ratings["movieId"]:
if x == 1:
print(y)

但是我的输出将永远保持打印，并迫使vs code关闭。

Answer 1

你应该在同一行上循环:

for x in zip(ratings["userId"], ratings["movieId"]):
if x == 1:
print(y)

或

res = ratings.loc[ratings["userId"].eq(1),"movieId"].tolist()

Answer 2

假设userId是字符串列，可以运行:

print(ratings.query("userId == '1'")["movieId"])

否则，如果是整数列，则可以运行:

print(ratings.query("userId == 1")["movieId"])