Skip to content

Instantly share code, notes, and snippets.

@neizod
Last active Sep 27, 2020
Embed
What would you like to do?
IMO 2020 Problem 3: Pebbles
#!/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)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment