Generic solution for rolling and then keeping arbitary number of (highest or lowest) dices

I'd like to calculate generic solution for expected value of dice roll NdX keep K where N = number of dices, X = number of die sides and K = number of dice to keep.

I have equations for special cases where N ≤ 2, K ≤ 1 or K ≥ N-1.

I just can't figure out non bruteforce generic solution.

Example python code [runnable]

#! /usr/bin/env python3

import random
import itertools

def avg_sim_keep_highest(n, d, keep):
return avg_sim(n, d, keep, 'highest')

def avg_sim_keep_lowest(n, d, keep):
return avg_sim(n, d, keep, 'lowest')

# simulate rolls, brute force if problem is simple enough
def avg_sim(n, d, keep, select):
N = 25000
if (d**n) < N:
return avg_brute(n, d, keep, select)
roll = lambda: sorted([random.randint(1, d) for _ in range(n)], reverse=(select=='highest'))[:keep]
return sum(sum(roll()) for _ in range(N)) / N

def avg_brute_keep_highest(n, d, keep):
return avg_brute(n, d, keep, 'highest')

def avg_brute_keep_lowest(n, d, keep):
return avg_brute(n, d, keep, 'lowest')

def avg_brute(n, d, keep, select):
die = range(1, d+1)
possible_rolls = itertools.product(die, repeat=n)
total_sum = sum(
sum(sorted(roll, reverse=(select=='highest'))[:keep])
for roll in possible_rolls
)

def avg_total(n, d):
return n*(d+1)/2

def avg_lowest(n, d):
return 1 + sum(x**n for x in range(d))/(d**n)

def avg_highest(n, d):
return d - sum(x**n for x in range(d))/(d**n)

def avg_drop_lowest(n, d):
return avg_total(n, d) - avg_lowest(n, d)

def avg_drop_highest(n, d):
return avg_total(n, d) - avg_highest(n, d)

def avg_2dX_lowest(n, d):
# (integrate n^2 from n=0 to X)/(X^2) + ((X^2-X)*(1/2) + X*(2/3))/(X^2)
assert n==2, "valid for 2 dice"
return ((1/3)*d**3 + (1/2)*d**2 + (1/6)*d) / (d**2)

def avg_2dX_highest(n, d):
return avg_total(n, d) - avg_2dX_lowest(n, d)

def avg_generic_highest(n, d, keep):
return "not implemented"

print("\nThrow 4d6 keep 3 highest")
print(avg_sim_keep_highest(n=4, d=6, keep=3))
print(avg_brute_keep_highest(n=4, d=6, keep=3))
print(avg_drop_lowest(n=4, d=6))

print("\nThrow 3d6 keep all")
print(avg_sim_keep_highest(n=3, d=6, keep=3))
print(avg_brute_keep_highest(n=3, d=6, keep=3))
print(avg_total(n=3, d=6))

print("\nThrow 4d6 keep 1 highest")
print(avg_sim_keep_highest(n=4, d=6, keep=1))
print(avg_brute_keep_highest(n=4, d=6, keep=1))
print(avg_highest(n=4, d=6))

print("\nThrow 2d20 keep 1 highest")
print(avg_sim_keep_highest(n=2, d=20, keep=1))
print(avg_brute_keep_highest(n=2, d=20, keep=1))
print(avg_2dX_highest(n=2, d=20))

print("\nThrow 2d20 keep 1 lowest")
print(avg_sim_keep_lowest(n=2, d=20, keep=1))
print(avg_brute_keep_lowest(n=2, d=20, keep=1))
print(avg_2dX_lowest(n=2, d=20))

print("\nThrow 4 keep 2 highest")
print(avg_sim_keep_highest(n=4, d=3, keep=2))
print(avg_brute_keep_highest(n=4, d=3, keep=2))
print(avg_generic_highest(n=4, d=6, keep=2))

Output

Throw 4d6 keep 3 highest
12.244598765432098
12.244598765432098
12.244598765432098

Throw 3d6 keep all
10.5
10.5
10.5

Throw 4d6 keep 1 highest
5.244598765432099
5.244598765432099
5.244598765432099

Throw 2d20 keep 1 highest
13.825
13.825
13.825

Throw 2d20 keep 1 lowest
7.175
7.175
7.175

Throw 4 keep 2 highest
5.08641975308642
5.08641975308642
not implemented