create pandas dataframe with random integers and finite sum across columns

Question

df = pd.DataFrame(np.random.randint(0,6,size=(1200000, 3)), 
columns=list('ABC'))
df['sum'] = df[['A','B','C']].sum(axis=1)
df = df[df['sum']==5]
df = df.sample(n=100000)

I want to create a dataframe with three different columns with random numbers between 0 and 5 such that sum across column is 5.

Answer 1

You can use itertools permutations to find numbers between 0 and 5 with sum equal to 5 and assign the result to DataFrame

import itertools
df = pd.DataFrame([elem for elem in list(itertools.permutations(range(6), 3)) if sum(elem) == 5], columns = list('ABC'))
df['sum'] = df.sum(1)

    A   B   C   sum
0   0   1   4   5
1   0   2   3   5
2   0   3   2   5
3   0   4   1   5
4   1   0   4   5
5   1   4   0   5
6   2   0   3   5
7   2   3   0   5
8   3   0   2   5
9   3   2   0   5
10  4   0   1   5
11  4   1   0   5