K-ary Symmetric Channel Capacity

Problem

Compute the capacity in bits of a k-ary symmetric channel with error probability p. With k >= 2 and 0 <= p <= 1, the formula is log2(k) + (1-p) log2(1-p) + p log2(p/(k-1)), and k = 2 reduces to the binary symmetric channel. Errors are uniform over the other k-1 symbols.

Hint

Start from log2(k) and subtract the two entropy-like correction terms.

Test cases

case 1

in: [2,0.1]

expected: 0.5310044064107188

case 2

in: [4,0.15]

expected: 1.1524153201754261

case 3

in: [3,0]

expected: 1.584962500721156

case 4

in: [2,0.5]

expected: 0

case 5

in: [7,1]

expected: 0.2223924213364481

Starter codePython

import math
def ksymmetric_capacity(k, p):
    # Your code here
    pass

5 test cases · run your code against them in the editor