Most efficient way to sum a conditional expression with numpy arrays

Viewed 37

I am currently working on a decision tree and need to calculate the entropy in a label vector y (an numpy array). I tried the two following methods to use the entropy formula:

entropy = sum(-len(y[y == label]) / len(y) * np.log2(len(y[y == label]) / len(y)) for label in labels)

and

        entropy = 0
        for label in labels:
            cond = y == label 
            fraction = len(y[cond]) / len(y)
            entropy += -fraction * np.log2(fraction)
        return entropy

When timing the two methods the seconds approach is the fastest. Why is that? I guess it is because I calculate the fraction two times in the same loop in the first approach. Is there any way to make the first approach faster?

"labels" in the above snippets are the following:

labels = np.unique(y)

1 Answers

You can even avoid the for loop by directly obtaining the count through np.unique:

def entropy(y):
    labels, count = np.unique(y, return_counts=True)
    prob = count / len(y)
    return -(prob * np.log2(prob)).sum()

Some performance comparisons with the for loop solution:

def for_loop_entropy(y):
    entropy = 0
    for label in np.unique(y):
        cond = y == label 
        fraction = len(y[cond]) / len(y)
        entropy += -fraction * np.log2(fraction)
    return entropy
In [1]: y = np.random.randint(0, 10, 2 * 10 ** 4)

In [2]: %timeit for_loop_entropy(y)
869 µs ± 9.17 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [3]: %timeit entropy(y)
369 µs ± 3.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [4]: y = np.random.randint(0, 100, 2 * 10 ** 4)

In [5]: %timeit for_loop_entropy(y)
2.58 ms ± 46.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [6]: %timeit entropy(y)
573 µs ± 5.84 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
Related