from itertools import *

# first two balls of same color
def first_two_balls(red, blue):
    total = 0
    good = 0
    for arrangement in combinations(range(red + blue), red):
        total = total + 1
        # if the first two balls are red, or they are both blue
        if (arrangement[0] == 0 and arrangement[1] == 1) or (arrangement[0] > 1):
            # this outcome is good
            good = good + 1
    return (good, total)


# choose two committees with no overlap
def committees(size, subsize):
    total = 0
    no_overlap = 0
    at_least_one = 0
    exactly_one = 0
    students = range(size)
    for c1 in combinations(students, subsize):
        for c2 in combinations(students, subsize):
            total = total + 1
            # calculate intersection size
            intsize = 0
            for x in c1:
                if x in c2:
                    intsize = intsize + 1
            if intsize == 0:
                no_overlap = no_overlap + 1
            else:
                at_least_one = at_least_one + 1
                if intsize == 1:
                    exactly_one = exactly_one + 1
    return (total, no_overlap, at_least_one, exactly_one)


# no two consecutive red balls
def no_consecutive_reds(red, blue):
    total = 0
    bad = 0
    for arrangement in combinations(range(red + blue), red):
        total = total + 1
        for i in range(red - 1):
            # if the i-th and (i+1)-st red balls are consecutive
            if arrangement[i + 1] - arrangement[i] == 1:
                # this outcome is bad
                bad = bad + 1
                break
    good = total - bad
    return good
    
