from z3 import *

# Partition range [1, N) into P sub-ranges of similar size

P = 4
N = Int('N')

sz = [Int('sz' + str(i)) for i in range(P)]

s = Solver()

# Precondition
s.add (N >= P)

# Implementation
for i in range(P):
  s.add (sz[i] == If(i < N%P,  N/P + 1, N/P))
  
# Check if "Add up to N" can be violated:
s.push()
s.add(Not(Sum([sz[i] for i in range(P)]) == N))
if s.check() == sat:
  print 'Subranges might not add up to N'
  print 'for N = ', s.model().evaluate(N)
  exit()
s.pop()
  

# Check if "similar" can be violated
for i in range(P):
  for j in range(P):
    s.push()
    s.add (Not(sz[i] - sz[j] <= 1))
    if s.check() == sat:
      print 'Subranges', i, 'and', j, 'might not be similar!'
      exit()
    s.pop()
    
print 'Verified for all N!'
                 
