IMO 2020 Problem 3: Pebbles

pebbles.py

#!/usr/bin/env python3

import sys
from argparse import ArgumentParser
from itertools import combinations
from random import seed, shuffle, randrange


sumsupto = lambda n: n*(n+1)//2
tco = lambda t: f'\033[01;{31+t[1]}m{t[0]:02}\033[00m'


parser = ArgumentParser(description='')
parser.add_argument('n', type=int, nargs='?', default=5)
parser.add_argument('r', type=int, nargs='?', default=42)
args = parser.parse_args()
N = args.n
R = args.r

seed(R)
cs = [i//4 for i in range(4*N)]
shuffle(cs)

xs = list(enumerate(cs, 1))
target = sumsupto(4*N) // 2

print('generated data:')
print(' '.join(tco(t) for t in xs))
print()

print('possible combinations:')
for ls in combinations(xs, 2*N):
    if ls[0][0] > 1: break          # hack for speed, rely on python combi order
    if any(ls[i][0]+ls[-i-1][0] != 4*N+1 for i in range(N)): continue       # only gaussian paring
    if sum(i for i, _ in ls) != target:
        continue
    if any(sum(1 for _, c in ls if c == i) != 2 for i in range(N)):
        continue
    rs = sorted(set(xs) - set(ls))
    print(' '.join(map(tco, ls)), '    ...    ', ' '.join(map(tco, rs)))