label = data[“label”]
pts = data[“jet_pt”]
top_idx = np.where(label[0,:] == 1)
top_pt = pts[top_idx]