检查pandas列是否仅包含0或1

import pandas as pd

data = {'col1': [0, 1, 0, 1, 0],
        'col2': [1, 1, 0, 1, 0],
        'col3': [0, 0, 0, 1, 1]}
df = pd.DataFrame(data)

def check_binary_column(column):
    unique_values = column.unique()
    if len(unique_values) == 2 and set(unique_values) == {0, 1}:
        return True
    else:
        return False

binary_columns = []
for column in df.columns:
    if check_binary_column(df[column]):
        binary_columns.append(column)

print("Binary columns:", binary_columns)