mjhjmtu

Given an input N as the size of a centrifuge, I need to find out the number of balanced configurations. The full description is here.

Centrifuge is a piece of equipment that puts multiple test tubes in rotation at very high speeds. It consists of a cylindrical rotor with holes situated evenly along the circumference of the rotor.
Because the device is operating at high speeds, the center of mass of all tubes must coincide with the center of the rotor. We assume all tubes have the same mass.

For the sake of simplicity, we also assume that tubes are point masses, fixed firmly to the rotor. All at the same distance from the center. You may safely assume the (x,y) coordinates of tube k are $Rsin a, Rcos a$, where $a = 2pifrackn$

The problem is: Given a centrifuge with N holes and K tubes, is it possible to balance it?

example 1: given N = 6 possible values for K are 2,3,4,6.

example 2: given N = 12, it is possible to balance 5 tubes: put 3 evenly dispersed (every 4th hole), then find a pair of opposite unoccupied holes and insert remaining 2 tubes there.

example 3: given N = 10, it is impossible to balance 7 tubes: put 5 evenly dispersed, and there's no opposite unoccupied holes left!

Input

Line 1: An integer N for the capacity of the centrifuge.

Output

Line 1: An integer M for the number of different possible values of K.

This is my code:

import numpy as np

N=int(input())
angle = 2*np.pi/N
z = complex(np.cos(angle), np.sin(angle))


import itertools
combs = 
lst = range(1,N)
for i in range(2, N-1):
 combs.append(i)
 els = [list(x) for x in itertools.combinations(lst, i)]
 combs.append(els)


count=1
lis = 
while count<=len(combs):
 for a in range(len(combs[count])):
 s=0
 for b in range(len(combs[count][a])):
 s+=z**combs[count][a][b]
 if abs(s)<1e-10:
 lis.append(combs[count][a])
 count+=2


lengths = 
for i in lis:
 if len(i) not in lengths:
 lengths.append(len(i))

print(len(lengths)+1)

The code works fine, but slows down after input size of above 20, and is practically unusable after the input grows to 50 or above, due to the for loops. How do I optimize it?

This is not a review, but an extended comment.

The code cannot be meaningfully optimized. Few notes though.

• First, thou shalt not bruteforce. It is almost inevitably wrong.




• Next, you must realize that it is a number-theoretical problem. Even the names of the test cases suggest so. Among other things it means that testing the sum against 1e-10 is not right.




• Finally, do watch the video referenced in the problem. It is a big spoiler, and it doesn't provide a proof of the claim. The proof is far from trivial; if you are interested in the proof, follow up with this is a fascinating reading (trigger warning: heavy-duty math).

There are two main problems:

1. Enumerating all combinations, even those that are obviously unbalanced




2. The use of floating point calculus in a combinatorics problem

Re-read the task carefully, especially examples 2 and 3, which hint a better algorithm:

1. Factorize N. The factors of N=12 are [2,3], which tells us that all tubes are members of either a balanced pair or triplet.




2. Find all combinations of factors that sum to K. For N=12, K=10, there are two solutions: 2+2+2+2+2=10 and 2+2+3+3=10.




3. Find out if the solutions are valid. For N=10, K=7 the only solution is 2+5=7, which is not valid (see example 3).

For large values of K, balance holes instead of tubes. So N=12, K=10 is the same as N=12, K=2.