Partition sieve algorithm for k-partitioning

partitionsieve.py

# ISC License
# Copyright (c) 2017, Lefteris Stamatogiannakis
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
# OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.
#
# partitionSieve function partitions an input list of numbers into multiple buckets. The max of any two bucket size differences is minimal.
#
# For more look here: https://plus.google.com/114866592715069940152/posts/QPwTJpn4he3

import random
from bisect import *
from itertools import *
from math import ceil

elems = [18235244038L, 29333785475L, 44241883158L, 90720783565L, 138278233069L, 160717749670L, 170810043347L, 190783321874L, 192761323717L, 263941209344L, 272732233974L, 283983798292L, 301164751968L,303174903185L, 316345714052L, 332247271357L, 354954626509L, 414848481662L, 477687829382L, 512677417648L, 535239618389L, 585259027380L, 620659086675L, 633354753524L, 651499393882L,656505134307L, 656531643207L, 709360592658L, 717553224165L, 729625007236L]

# elems = [random.randint(0, 2 ** 32) for _ in range(30)]

def checkSolution(s):
    sums = [sum(i) for i in s]
    return max(sums), min(sums)

def kk(elems, k):
    k1 = k - 1
    e = [(tuple([x] + [0] * (k - 1)), [[x]] + [[] for x in range(k1)]) for x in sorted(elems)]
    while len(e) > 1:
        e1 = e.pop()
        e2 = e.pop()
        edif = sorted([e1[0][i] + e2[0][k1-i] for i in range(k)], reverse = True)
        edif = tuple([x - edif[-1] for x in edif])
        comb = sorted([e1[1][i] + e2[1][k1-i] for i in range(k)], key = sum, reverse = True)
        insort(e, (edif, comb))
    return e[0][1]

def partitionSieve(elems, k, highcount = None, memsize = 23.8):
    elemlen = len(elems)
    elemlen1 = elemlen - 1
    elems = sorted(elems)
    lowstart = 0
    totalsum = sum(elems)
    c = kk(elems, k)
    smallest = c
    k1 = k - 1
    ub, _ = checkSolution(c)
    ub1 = ub - 1
    lb = totalsum - k1 * ub1

    if highcount is None:
        frontsum = 0
        frontrange = 0
        for i in xrange(elemlen):
            frontsum += elems[i]
            if frontsum >= ub:
                frontrange = i
                break

        backsum = 0
        backrange = 0
        for i in xrange(elemlen-1, -1, -1):
            backsum += elems[i]
            if backsum >= ub:
                backrange = elemlen - 1 - i
                break
        if backrange == 1:
            highcount = elemlen
        else:
            highcount = min(int(2.0**(elemlen/4.5) * ((2.0**float(backrange))/(2.0**float(frontrange))) * elemlen), elemlen)

    elemshigh = sorted(elems[elemlen - highcount:], reverse = True)
    elems = elemshigh + elems[:elemlen - highcount]
    lowstart = len(elemshigh)

    allcomb = [None] * elemlen
    totalcomb = 0
    combi = 0
    allcomb[-1] = [0, elems[-1]]
    comb = set(allcomb[-1])
    combs = ub
    for i in xrange(elemlen - 2, -1, -1):
        curelem = elems[i]
        for x in allcomb[i + 1]:
            nelem = curelem + x
            if nelem >= ub: break
            comb.add(nelem)
        allcomb[i] = list(comb)
        allcomb[i].sort()
        totalcomb += len(allcomb[i])
        combi = i
        combs = -1 if i > lowstart else (ub if i ==0 else elems[i-1])
        if totalcomb > int(2 ** memsize):
            break

    comb = None
    optb = int(ceil(float(totalsum) / k))
    stack = []
    stack.append(([elems[0]], [[elems[0]]], 1))
    disp0 = [0] * k
    lstack = []
    nodevisit = 0
    childslived = 0
    childsvisited = 0
    leafvisit = 0
    seenelems = set()
    levelcount = [0] * elemlen
    while True:
        if len(lstack) > 0:
            sums, c, s = lstack.pop()
            if len(lstack) > 0:
                stack += lstack
                del(lstack[:])
        else:
            if len(stack) == 0:
                break
            sums, c, s = stack.pop()

        nodevisit += 1

        if len(sums) < k:
            c.append([])
            sums.append(0)

        sp1 = s + 1
        curelem = elems[s]

        newdisp = disp0[:]
        newdispsum = 0
        if sp1 < elemlen:
            for i in xrange(k):
                nsum = sums[i] if i < len(sums) else 0
                ubd = ub1 - nsum
                if sp1 >= lowstart:
                    if nsum + elems[sp1] > ub1:
                        d = ub1 - nsum
                    elif sp1 >= combi:
                        mcs = max(sp1, combi)
                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                        d = ub1 - lub
                    else:
                        d = 0
                else:
                    if ubd < combs or sp1 >= combi:
                        mcs = max(sp1, combi)
                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                        d = ub1 - lub
                    else:
                        d = 0

                newdisp[i] = d
                newdispsum += d

        stopcond = False
        childsvisited += len(sums)
        seenelems.clear()
        for i in xrange(len(sums)):
            if stopcond: break
            if sums[i] in seenelems:
                continue
            else:
                seenelems.add(sums[i])

            nsum = sums[i] + curelem

            if nsum > ub1:
                continue

            if s < elemlen1:
                skiplub = ub1 - newdisp[i]
                if sp1 < lowstart and nsum >= skiplub:
                    del(lstack[:])
                    stopcond = True

                if lb > skiplub:
                    del(lstack[:])
                    stopcond = True

                llb = lb + newdispsum - newdisp[i]
                if llb < lb:
                    continue

                ubd = ub1 - nsum
                if sp1 >= lowstart:
                    if nsum + elems[sp1] > ub1:
                        lub = nsum
                    elif sp1 >= combi:
                        mcs = max(sp1, combi)
                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                    else:
                        lub = ub1

                    if llb > lub:
                        continue
                else:
                    if ubd < combs or sp1 >= combi:
                        mcs = max(sp1, combi)
                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                        if llb > lub:
                            continue

            newcomb = c[:]
            newcomb[i] = newcomb[i][:]
            newcomb[i].append( curelem )
            newsums = sums[:]
            newsums[i] = nsum

            if s < elemlen1:
                levelcount[s] += 1
                lstack.append((newsums, newcomb, sp1))
            else:
                leafvisit += 1
                if max(newsums) < ub:
                    smallest = newcomb[:]
                    ub = max(newsums)
                    ub1 = ub - 1
                    lb = totalsum - k1 * ub1
                    newstack = []
                    dedup = set()
                    for x in stack:
                        lsums = x[0]
                        if max(lsums) >= ub: continue

                        sig = tuple(sorted(lsums))
                        if sig in dedup:
                            continue
                        else:
                            dedup.add(sig)

                        lsp1 = x[2]
                        lnewdisp = [0] * k
                        if lsp1 < elemlen:
                            for i1 in xrange(k):
                                nsum = lsums[i1] if i1 < len(lsums) else 0
                                ubd = ub1 - nsum
                                if lsp1 >= lowstart:
                                    if nsum + elems[lsp1] > ub1:
                                        d = ub1 - nsum
                                    elif lsp1 >= combi:
                                        mcs = max(lsp1, combi)
                                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                                        d = ub1 - lub
                                    else:
                                        d = 0
                                else:
                                    if ubd < combs or lsp1 >= combi:
                                        mcs = max(lsp1, combi)
                                        lub = nsum + allcomb[mcs][bisect(allcomb[mcs], ubd) - 1]
                                        d = ub1 - lub
                                    else:
                                        d = 0

                                lnewdisp[i1] = d

                            lasti = 0
                            maxllb = 0
                            for i1 in xrange(k):
                                maxllb = max(maxllb, lb + sum(lnewdisp) - lnewdisp[i1])

                            if maxllb > ub1 or maxllb < lb:
                                continue

                        newstack.append(((maxllb, max(lsums), sum(lsums)), x[0], x[1], x[2]))
                    newstack.sort(reverse = True)
                    stack = [(x[1],x[2],x[3]) for x in newstack]
                    if ub == optb:
                        return c

        childslived += max(len(lstack) - 1, 0)
    return smallest

print "Partitioning", len(elems),"items:"
print elems
print

k = 10
print "into", k, "buckets:"
print partitionSieve(elems, k)