HITCON QUALS 2016 - Beelzemon (PPC 150)

beelzemon_solver.py

#!/usr/bin/env python
#-*- coding:utf-8 -*-

from sock import Sock

def find_greedy(s):
    a, b = [], []
    sa = sb = 0
    for x in sorted(s, reverse=True, key=abs):
        if abs(sa + x - sb) < abs(sb + x - sa):
            a.append(x)
            sa += x
        else:
            b.append(x)
            sb += x
    while len(a) < len(b):
        a.append(b.pop())
    while len(a) > len(b):
        b.append(a.pop())
    if sum(a) == sum(b):
        return a
    return find_random(s, a, b)

def find_swap(a, b, sa, sb):
    i = j = 0
    best = (0, 0)
    best_sa = -999999999999999
    wanted_sa = (sa + sb) / 2
    while i < len(a) and j < len(b):
        test_sa = sa + b[j] - a[i]
        if abs(test_sa - wanted_sa) < abs(best_sa - wanted_sa):
            best = i, j
            best_sa = test_sa
        if test_sa < wanted_sa:
            j += 1
        elif test_sa > wanted_sa:
            i += 1
        elif test_sa == wanted_sa:
            if i < len(a) - 1:
                i += 1
            else:
                j += 1
    return best

def move(a, i):
    while i > 0 and a[i] < a[i-1]:
        a[i-1], a[i] = a[i], a[i-1]
        i -= 1
    while i < len(a) - 1 and a[i] > a[i+1]:
        a[i+1], a[i] = a[i], a[i+1]
        i += 1

def find_random(s, a, b):
    wanted = sum(s) / 2
    sa, sb = sum(a), sum(b)
    a = sorted(a)
    b = sorted(b)
    bad_itr = 0
    prev = 0
    while 1:
        i, j = find_swap(a, b, sa, sb)
        sa = sa - a[i] + b[j]
        sb = sb - b[j] + a[i]
        a[i], b[j] = b[j], a[i]
        move(a, i)
        move(b, j)
        print n, k, len(s), "sa", sa, wanted, abs(sa - wanted)
        if sa == wanted:
            return a

        if abs(sa - wanted) >= prev:
            bad_itr += 1
            if bad_itr == 10:
                print "----"
                break
        else:
            bad_itr = 0
            prev = abs(sa - wanted)

f = Sock("52.198.217.117 6666")
itr = 0
while 1:
    itr += 1
    try:
        res = f.read_line()
        n, k = map(int, res.split())
    except Exception as err:
        print "Flag?", err, res, `f.buf`
        quit()

    print "#%d: %d %d" % (itr, n, k)

    s = range(-2**n, 2**n)
    if k & 1:
        invmap = {x**k: x for x in s}
    else:
        invmap = {x**k: -abs(x) for x in s}
    sk = sorted([x**k for x in s], key=abs)

    print "solving.."
    res = find_greedy(sk)
    assert res, "failed :("
    ans = []
    for x in res:
        ans.append(invmap[x])
        if k % 2 == 0:
            invmap[x] *= -1
    assert sum(x**k for x in ans) == sum(sk) / 2

    vec = " ".join(map(str, ans))
    f.send_line(vec)


for n in xrange(1, 21):
    for k in xrange(1, n + 1):
        s = range(-2**n, 2**n)
        sk = sorted([x**k for x in s],key=abs)
        if k & 1:
            invmap = {x**k: x for x in s}
        else:
            invmap = {x**k: -abs(x) for x in s}

        print n, k
        res = find_greedy(sk)
        assert res, "failed :("
        ans = []
        for x in res:
            ans.append(invmap[x])
            if k % 2 == 0:
                invmap[x] *= -1
        assert sum(x**k for x in ans) == sum(sk) / 2