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

# 
# <div style="text-align:right">Peter Norvig<br>Nov 2019</div>
# 
# # Riddler Lottery
# 
# The 538 Riddler [poses](https://fivethirtyeight.com/features/can-you-decode-the-riddler-lottery/) this problem:
# 
# > Five friends with a lot in common are playing the [Riddler Lottery](https://fivethirtyeight.com/features/can-you-decode-the-riddler-lottery/), in which each must choose exactly five numbers from 1 to 70. After they all picked their numbers,
# - The first friend notices that no number was selected by two or more friends. 
# - The second friend observes that all 25 selected numbers are composite (i.e., not prime). 
# - The third friend points out that each selected number has at least two distinct prime factors. 
# - The fourth friend excitedly remarks that the product of selected numbers on each ticket is exactly the same. 
# - The fifth friend is left speechless. (You can tell why all these people are friends.)
# 
# > 1. What is the product of the selected numbers on each ticket?
# 2. How many _different_ ways could the friends have selected five numbers each so that all their statements are true?
# 
# # Preliminary Analysis
# 
# The fourth friend's statement was a bit unclear, but I take it to mean that each friend multiplied together their own five numbers, and they all got the same product.  To be concrete, here's an example of a solution in a simplified version of the problem where each friend only selects two tickets, not five:
# 
#      Friend   Selection  Product  Factors
#      1        ( 6, 60)   360      [2, 3] + [2, 2, 3, 5]
#      2        (10, 36)   360      [2, 5] + [2, 2, 3, 3]
#      3        (12, 30)   360      [2, 2, 3] + [2, 3, 5]
#      4        (15, 24)   360      [3, 5] + [2, 2, 2, 3]
#      5        (18, 20)   360      [2, 3, 3] + [2, 2, 5]
# 
# And here's a list of the key concepts:
# 
# - **number**: An integer from 1 to 70, e.g. the int `42`.
# - **factors**: Every positive integer has a unique prime factorization, e.g. `factors(12) == [2, 2, 3]`:  two distinct primes factors. But `factors(8) == [2, 2, 2]`: one distinct prime factor.
# - **selection**: A collection of 5 numbers, e.g. the sorted tuple `(12, 15, 20, 28, 30)`.
# - **product**: The result of multiplying together the 5 numbers in a selection, e.g. the int `3024000`.
# - **candidate**: A candidate solution is a set of 5 selections e.g.  `{( 6, 60), (10, 36), (12, 30), (15, 24), (18, 20)}` in my simplified version where each selection has only two numbers.
# - **solution**: A solution is a candidate that satisifes each of the four friends' statements.
# 
# Can I use brute force and enumerate all the possible candidates? 
# 
# There are (70 choose 5) × (65 choose 5) × (60 choose 5) × (55 choose 5) × (50 choose 5) / 5! or [about](https://www.wolframalpha.com/input/?i=%2870+choose+5%29+*+%2865+choose+5%29+*+%2860+choose+5%29+*+%2855+choose+5%29+*+%2850+choose+5%29+%2F+5%21) $10^{31}$ candidates, so no. We'll have to be more clever.

# # Valid Numbers
# 
# There will be fewer candidates to consider if we can reduce the number of valid numbers to select from. The __third__ friend stated that the numbers all have at least two distinct prime factors. So let's find the numbers that have that property. (The numbers we are dealing with are small, so don't worry about the inefficiency of my function `factors`.)

# In[1]:


from itertools import combinations
from collections import Counter, defaultdict


# In[2]:


def factors(n) -> list:
    "List of prime factors that multiply together to give n."
    return ([]   if n == 1
            else next([p] + factors(n // p) 
                      for p in range(2, n + 1) if n % p == 0))

distinct = set

numbers = {n for n in range(1, 71) if len(distinct(factors(n))) >= 2}

print(len(numbers), numbers)


# Great; we got it down from 70 to 41 possible numbers.
# 
# Now the __fourth__ friend's statement says that each friend picks five numbers that have the same product. 
# In my simplified version where the friends pick two numbers each, they all picked a selection with the product 360.  The prime factorization of 360 is `[2, 2, 2, 3, 3, 5]`; that means that all five friends had to find a different way of allocating these factors to their numbers. 
# 
# For each number that is selected by any friend, and for each prime factor $p$ of that number, it must be the case that there are at least four other valid numbers that also have $p$ as a factor; otherwise the product couldn't be the same for all the friends. So let's count in how many numbers each prime appears:

# In[3]:


prime_counts = Counter(p for n in numbers for p in distinct(factors(n)))

prime_counts


# This says that the prime factor 2 appears in 29 valid numbers and the prime factor 31 appears in only 1 valid number. Only factors that appear in at least 5 valid numbers can be part of a solution, so that's `{2, 3, 5, 7, 11}`.
# 
# Let's update the set of valid numbers to contain only numbers $n$ such that every prime factor $p$ of $n$ appears in at least 5 valid numbers:

# In[4]:


numbers = {n for n in numbers if all(prime_counts[p] >= 5 for p in distinct(factors(n)))}

print(len(numbers), numbers)


# There are now only 28 valid numbers; a nice reduction from 70 to 41 to 28.
# 
# # Valid Selections
# 
# Now let's switch attention from individual numbers to selections of five numbers. There are (28 choose 5) = 98,280 possible selections; a manageable number. But that means there are (98,280 choose 5) = $\approx 10^{23}$ candidate solutions; an unmanageable number.
# 
# My first thought to reduce the number of candidates is to say that we should only consider candidates where all five selections in the candidate have the same product. To do that, we can group selections by product.
# 
# We'll make `products` be a `dict` where each key is the product of the five numbers in a selection, and the corresponding value is a list of all the selections of five numbers with that product:

# In[5]:


def multimap(items, key) -> dict:
    "A dict of {key(item): [item, ...]}"
    result = defaultdict(list)
    for x in items:
        result[key(x)].append(x)
    return result

def product(nums) -> int: 
    "Multiply nums together (similar to sum(nums))."
    result = 1
    for num in nums:
        result *= num
    return result

products = multimap(combinations(numbers, 5), key=product)


# Here is an entry in the `products` dict:

# In[6]:


products[6 * 12 * 14 * 15 * 20]


# This says there are 6 selections whose product is 6 * 12 * 14 * 15 * 20 = 302,400; if there were a way to choose 5 of these 6 with all distinct numbers, that would be a solution. Sadly, there is no such way. Since every one of the selections contains a 6; we can't even choose two disjoint selections, let alone five.
# 
# Let's see how many different products there are, and how many have at least five selections:

# In[7]:


len(products), len([n for n in products if len(products[n]) >= 5])


# It seems reasonable to go through all the products and see if one of the 29,506 with 5 or more selections can come up with 5 disjoint selections. The function `k_disjoint(k, selections)` finds all ways to choose `k` different elements of `selections` such that there is no shared number among any selection. The function keeps track of a `partial_solution`&mdash;a set of previously-found selections&mdash;as it recursively searches for a complete solution. Any new selection must be disjoint from all the selections in `partial_solution`. 

# In[8]:


def k_disjoint(k, selections, start=0, partial_solution=set()) -> list:
    "All ways of picking k elements of selections that have all disjoint members."
    if len(partial_solution) == k:
        yield partial_solution
    elif len(partial_solution) + (len(selections) - start) >= k:
        for i in range(start, len(selections)):
            selection = selections[i]
            if all(is_disjoint(selection, s) for s in partial_solution):
                yield from k_disjoint(k, selections, i + 1, partial_solution | {selection})
    
def is_disjoint(A, B) -> bool: return not any(a in B for a in A)


# Here's an example of how `k_disjoint` works. Out of the following six selections (which in this simplified example have only two numbers each, not five), there are two ways to pick five selections without having a duplicate number:

# In[9]:


selections = [(6, 60), (10, 36), (12, 30), (15, 24), (18, 20), (10, 35)]
list(k_disjoint(5, selections))


# But for the six selections whose product is 302,400, there are no disjoint solutions:

# In[10]:


selections = [(6, 10, 12, 14, 30),
 (6, 10, 12, 15, 28),
 (6, 10, 12, 20, 21),
 (6, 10, 14, 15, 24),
 (6, 10, 14, 18, 20),
 (6, 12, 14, 15, 20)]
list(k_disjoint(5, selections))


# Now we're ready to solve the problem.
# 
# # 1. What is the product of the selected numbers on each ticket?
# 
# That is, find the  number `N` in `products` that can form 5 disjoint selections.

# In[11]:


N = next(n for n in products if any(k_disjoint(5, products[n])))

print(f'The product is {N:,d}; factors are {factors(N)}')


# # 2. How many different ways could the friends have selected five numbers?
# 
# I'll compute all of the results for `k_disjoint`, and see how many there are:

# In[12]:


different_ways = list(k_disjoint(5, products[N]))
print(f'There are {len(different_ways):,d} different ways.')


# That's too many to look at all of them, but I can peek at every thousandth one:

# In[13]:


different_ways[::1000]