Python - ссылка на столбец dataframe с использованием str () вместо кавычек

Question

Мой код - это функция, которая вычисляет конъюнктивную вероятность флагов стран в образце, содержащем определенный цвет.Это из упражнения DataQuest.Он запрашивает только вероятность трех флагов, содержащих красный, но я хотел бы бросить себе вызов и написать функцию для n и цвета.

Flags - это фрейм данных.Есть цветные столбцы.1, если цвет есть, 0, если его нет.

import numpy as np

def conjunctive_probability(n, colour):
    total_count = flags.shape[0]
    colour_picked = flags[flags[str(colour)] == 1].shape[0]
    p = 0
    probabilities = []

    for p in range(n):
        probability = colour_picked / total_count
        probabilities.append(probability)
        colour_picked -= 1
        total_count -= 1
        p += 1
    return np.prod(np.array(probabilities))

three_red = conjunctive_probability(3, red)

В строке 5 появляется ошибка (colour_picked).Если я наберу там цвет, например:

colour_picked = flags[flags['red'] == 1].shape[0]

, это сработает.

Но я не понимаю, почему str () не работает.Это дает мне:

KeyError: '153'

, которое является количеством флагов, имеющих красный цвет.

Answer 1

str() не проблема.Вероятно, ключевая ошибка указывает на то, что flags[YOUR_KEY] не существует.В конкретной ошибке flags[153] не существует, что означает, что 153 не является ключом в вашем flags словечке.Я не вижу, где flags инициализируется в вашем посте.