Sum Balance (Codeforces 1242C)
Tags: DP, bitmasks

Ujan has a lot of numbers in his boxes. He likes order and balance, so he decided to reorder the numbers.There are $$$k$$$ boxes numbered from $$$1$$$ to $$$k$$$. The $$$i$$$-th box contains $$$n_i$$$ integer numbers. The integers can be negative. All of the integers are distinct. Ujan is lazy, so he will do the following reordering of the numbers exactly once. He will pick a single integer from each of the boxes, $$$k$$$ integers in total. Then he will insert the chosen numbers — one integer in each of the boxes, so that the number of integers in each box is the same as in the beginning. Note that he may also insert an integer he picked from a box back into the same box.Ujan will be happy if the sum of the integers in each box is the same. Can he achieve this and make the boxes perfectly balanced, like all things should be?

Input: The first line contains a single integer $$$k$$$ ($$$1 \leq k \leq 15$$$), the number of boxes. The $$$i$$$-th of the next $$$k$$$ lines first contains a single integer $$$n_i$$$ ($$$1 \leq n_i \leq 5\,000$$$), the number of integers in box $$$i$$$. Then the same line contains $$$n_i$$$ integers $$$a_{i,1}, \ldots, a_{i,n_i}$$$ ($$$|a_{i,j}| \leq 10^9$$$), the integers in the $$$i$$$-th box. It is guaranteed that all $$$a_{i,j}$$$ are distinct.

Output: If Ujan cannot achieve his goal, output "No" in a single line. Otherwise in the first line output "Yes", and then output $$$k$$$ lines. The $$$i$$$-th of these lines should contain two integers $$$c_i$$$ and $$$p_i$$$. This means that Ujan should pick the integer $$$c_i$$$ from the $$$i$$$-th box and place it in the $$$p_i$$$-th box afterwards.If there are multiple solutions, output any of those.You can print each letter in any case (upper or lower).

Note: In the first sample, Ujan can put the number $$$7$$$ in the $$$2$$$nd box, the number $$$2$$$ in the $$$3$$$rd box, the number $$$5$$$ in the $$$1$$$st box and keep the number $$$10$$$ in the same $$$4$$$th box. Then the boxes will contain numbers $$$\{1,5,4\}$$$, $$$\{3, 7\}$$$, $$$\{8,2\}$$$ and $$$\{10\}$$$. The sum in each box then is equal to $$$10$$$.In the second sample, it is not possible to pick and redistribute the numbers in the required way.In the third sample, one can swap the numbers $$$-20$$$ and $$$-10$$$, making the sum in each box equal to $$$-10$$$.

Sample Inputs:
4
3 1 7 4
2 3 2
2 8 5
1 10

2
2 3 -2
2 -1 5

2
2 -10 10
2 0 -20

Sample Outputs:
Yes
7 2
2 3
5 1
10 4

No

Yes
-10 2
-20 1

Solution (Python):
from itertools import accumulate
from sys import stdin, stdout

def main():
    k = int(stdin.readline())
    a = [tuple(map(int, stdin.readline().split()[1:])) for _ in range(k)]
    a2ij = {aij: (i, j) for (i, ai) in enumerate(a) for (j, aij) in enumerate(ai)}
    plena = [0] + list(accumulate(map(len, a)))
    suma = tuple(map(sum, a))
    totala = sum(suma)
    if totala % k != 0:
        stdout.write('No\n')
    else:
        needle = totala // k
        mask2i2cp = compute_mask2i2cp(a, a2ij, needle, plena, suma)
        dp = compute_previous_mask(mask2i2cp)
        output(dp, mask2i2cp)

def compute_mask2i2cp(a, a2ij, needle, plena, suma):
    used = [False] * plena[-1]
    number_of_masks = 1 << len(a)
    mask2i2cp = [-1] * number_of_masks
    for (i, ai) in enumerate(a):
        for (j, aij) in enumerate(ai):
            if not used[plena[i] + j]:
                (mask, i2cp) = compute_mask_i2cp(a2ij, aij, i, j, needle, suma)
                if i2cp != -1:
                    mask2i2cp[mask] = i2cp
                    for cp in i2cp:
                        if cp != -1:
                            (c, p) = cp
                            (ii, jj) = a2ij[c]
                            used[plena[ii] + jj] = True
    return mask2i2cp

def output(dp, mask2i2cp):
    mask = len(mask2i2cp) - 1
    if dp[mask] == -1:
        stdout.write('No\n')
    else:
        answer = [-1] * len(mask2i2cp[dp[mask]])
        while mask > 0:
            current_mask = dp[mask]
            for (i, cp) in enumerate(mask2i2cp[current_mask]):
                if 1 == current_mask >> i & 1:
                    (c, p) = cp
                    answer[i] = (c, p)
            mask ^= current_mask
        stdout.write('Yes\n' + '\n'.join(('{} {}'.format(c, 1 + p) for (c, p) in answer)))

def compute_mask_i2cp(a2ij, aij, i, j, needle, suma):
    i2cp = [-1] * len(suma)
    mask = 0
    current_a = aij
    current_i = i
    try:
        while True:
            next_a = needle - (suma[current_i] - current_a)
            (next_i, next_j) = a2ij[next_a]
            if mask >> next_i & 1 == 1:
                return (mask, -1)
            mask |= 1 << next_i
            i2cp[next_i] = (next_a, current_i)
            if next_i == i:
                if next_j == j:
                    return (mask, i2cp)
                return (mask, -1)
            if next_i == current_i:
                return (mask, -1)
            current_a = next_a
            current_i = next_i
    except KeyError:
        return (mask, -1)

def compute_previous_mask(mask2cp):
    number_of_masks = len(mask2cp)
    dp = [-1] * number_of_masks
    dp[0] = 0
    for (mask, cp) in enumerate(mask2cp):
        if cp != -1:
            complement_mask = number_of_masks - 1 & ~mask
            previous_mask = complement_mask
            while previous_mask > 0:
                if dp[previous_mask] != -1 and dp[previous_mask | mask] == -1:
                    dp[previous_mask | mask] = mask
                previous_mask = previous_mask - 1 & complement_mask
            if dp[mask] == -1:
                dp[mask] = mask
    return dp
if __name__ == '__main__':
    main()

Solution (Python):
from itertools import accumulate
from sys import stdin, stdout

def main():
    k = int(stdin.readline())
    a = [tuple(map(int, stdin.readline().split()[1:])) for _ in range(k)]
    a2ij = {aij: (i, j) for (i, ai) in enumerate(a) for (j, aij) in enumerate(ai)}
    plena = [0] + list(accumulate(map(len, a)))
    suma = tuple(map(sum, a))
    totala = sum(suma)
    if totala % k != 0:
        stdout.write('No\n')
    else:
        needle = totala // k
        mask2i2cp = compute_mask2i2cp(a, a2ij, needle, plena, suma)
        dp = compute_previous_mask(mask2i2cp)
        output(dp, mask2i2cp)

def compute_mask2i2cp(a, a2ij, needle, plena, suma):
    used = [False] * plena[-1]
    number_of_masks = 1 << len(a)
    mask2i2cp = [-1] * number_of_masks
    for (i, ai) in enumerate(a):
        for (j, aij) in enumerate(ai):
            if not used[plena[i] + j]:
                (mask, i2cp) = compute_mask_i2cp(a2ij, aij, i, j, needle, suma)
                if i2cp != -1:
                    mask2i2cp[mask] = i2cp
    return mask2i2cp

def output(dp, mask2i2cp):
    mask = len(mask2i2cp) - 1
    if dp[mask] == -1:
        stdout.write('No\n')
    else:
        answer = [-1] * len(mask2i2cp[dp[mask]])
        while mask > 0:
            current_mask = dp[mask]
            for (i, cp) in enumerate(mask2i2cp[current_mask]):
                if 1 == current_mask >> i & 1:
                    (c, p) = cp
                    answer[i] = (c, p)
            mask ^= current_mask
        stdout.write('Yes\n' + '\n'.join(('{} {}'.format(c, 1 + p) for (c, p) in answer)))

def compute_mask_i2cp(a2ij, aij, i, j, needle, suma):
    i2cp = [-1] * len(suma)
    mask = 0
    current_a = aij
    current_i = i
    try:
        while True:
            next_a = needle - (suma[current_i] - current_a)
            (next_i, next_j) = a2ij[next_a]
            if mask >> next_i & 1 == 1:
                return (mask, -1)
            mask |= 1 << next_i
            i2cp[next_i] = (next_a, current_i)
            if next_i == i:
                if next_j == j:
                    return (mask, i2cp)
                return (mask, -1)
            if next_i == current_i:
                return (mask, -1)
            current_a = next_a
            current_i = next_i
    except KeyError:
        return (mask, -1)

def compute_previous_mask(mask2cp):
    number_of_masks = len(mask2cp)
    dp = [-1] * number_of_masks
    dp[0] = 0
    for (mask, cp) in enumerate(mask2cp):
        if cp != -1:
            complement_mask = number_of_masks - 1 & ~mask
            previous_mask = complement_mask
            while previous_mask > 0:
                if dp[previous_mask] != -1 and dp[previous_mask | mask] == -1:
                    dp[previous_mask | mask] = mask
                previous_mask = previous_mask - 1 & complement_mask
            if dp[mask] == -1:
                dp[mask] = mask
    return dp
if __name__ == '__main__':
    main()

Solution (Python):
from collections import defaultdict
data = defaultdict(list)
position = defaultdict()
nxt = defaultdict()
agg_sum = list()
k = int(input())
trace = defaultdict()
F = [False for x in range(1 << k)]
back = [0 for x in range(1 << k)]
total_sum = 0
res = [(0, 0) for x in range(k)]

def build_mask(trace_mask):
    if trace_mask == 0:
        return
    if trace.get(trace_mask):
        for data in trace.get(trace_mask):
            (fr, to, v) = data
            res[fr] = (v, to)
        return
    sub_mask = back[trace_mask]
    build_mask(sub_mask)
    build_mask(trace_mask - sub_mask)
if __name__ == '__main__':
    for i in range(k):
        values = list(map(int, input().split(' ')))
        data[i] = values[1:]
        agg_sum.append(sum(data[i]))
        total_sum += agg_sum[i]
        for (cnt, v) in enumerate(data[i], 0):
            position[v] = (i, cnt)
    if total_sum % k != 0:
        print('No')
        exit(0)
    row_sum = total_sum // k
    for i in range(k):
        for (cnt, value) in enumerate(data.get(i), 0):
            x = i
            y = cnt
            mask = 1 << x
            could = True
            circle = list()
            while True:
                next_value = row_sum - agg_sum[x] + data.get(x)[y]
                if position.get(next_value) is None:
                    could = False
                    break
                last_x = x
                last_y = y
                (x, y) = position.get(next_value)
                circle.append((x, last_x, next_value))
                if x == i and y == cnt:
                    break
                if mask & 1 << x:
                    could = False
                    break
                mask |= 1 << x
            F[mask] |= could
            if could:
                trace[mask] = circle
    for mask in range(1, 1 << k):
        sub = mask
        while sub > 0:
            if F[sub] and F[mask - sub]:
                F[mask] = True
                back[mask] = sub
                break
            sub = mask & sub - 1
    if F[(1 << k) - 1]:
        print('Yes')
        build_mask((1 << k) - 1)
        for value in res:
            print(value[0], value[1] + 1)
    else:
        print('No')