I am trying to randomly generate integer partitions (N into M parts) in pytorch with a minimum partition size of 1.
For example, (3, 1, 1)
and (4, 1, 0)
are both partitions of 5 into 3 parts but (4, 1, 0)
's minimum partition size is 0 so this should not be allowed
I would like to use this to generate my dataset on demand, so would nice if there was a pytorch (parrallel/gpgpu) solution.
See other questions about generating integer partitions:
- Generate restricted weak integer compositions (or partitions) of an integer n into k parts in Python
- Elegant Python code for Integer Partitioning
CodePudding user response:
Here's a solution that works in limited cases, M hast to divide N, 2 has to divide M, and the maximium is limited, but this is the behaviour I wanted.
You start of with the equal partition then calculate a delta that sums to zero...
M = 4
N = 16
MINIMUM = 1
assert N % M == 0
assert M % 2 == 0
avg = N // M
equal_partition = torch.full((M,), avg)
half_delta = torch.randint(-(avg-MINIMUM), avg-MINIMUM, (M // 2,))
delta = torch.cat((half_delta, -half_delta), dim=-1)[torch.randperm(M)]
partition = equal_partition delta
print(partition, partition.sum())