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

# <div style="text-align:right"><i>Peter Norvig<br>Nov 2019<br>revised May 2021</i></div>
# 
# # Riddler Lottery
# 
# The 538 Riddler [poses](https://fivethirtyeight.com/features/can-you-decode-the-riddler-lottery/) this problem:
# 
# > Five friends are playing the [Riddler Lottery](https://fivethirtyeight.com/features/can-you-decode-the-riddler-lottery/), in which each player selects exactly five integers from 1 to 70. After they all picked their numbers,
# 1. The first player states that no number was selected by two or more players. 
# 2. The second player observes that all the selected numbers are composite (i.e., not prime). 
# 3. The third player points out that each selected number has at least two distinct prime factors. 
# 4. The fourth player notes that each player's 5 numbers multiply to the same product. 
# 5. The fifth player is left speechless. 
# >
# > That leads to two questions:
# > - What is the unique product, $P$, that each player's 5 numbers multiply together to?
# > - In how many different ways could the selections be made so that the statements above are true?
# 
# #  Analysis
# 
# As an example, consider a version of the problem where each player selects 2 numbers, not 5. Here is one solution:
# 
#       
# |Player |Selection| Product |Prime Factors|
# |-|-|-|-|
# |1  |   6, 60 | 360  |   {2, 3}, {2, 3, 5}
# |2  |  10, 36 | 360  |   {2, 5}, {2, 3}
# |3  |  12, 30 | 360  |   {2, 3}, {2, 3, 5}
# |4  |  15, 24 | 360  |   {3, 5}, {2, 3}
# |5  |  18, 20 | 360  |   {2, 3}, {2, 5}
# 
# 
# 
# The key concepts:
# 
# - **Numbers**: The integers from 1 to 70. 
# - **Selection**: A sorted tuple of numbers, e.g. `(6, 60)` for 2 numbers, or `(12, 15, 20, 28, 30)` for 5.
# - **Candidate**: A set of 5 selections e.g. `{(6, 60), (10, 36), (12, 30), (15, 24), (18, 21)}`.
# - **Solution**: A candidate that satisfies statements 1–4.
# - **Distinct prime factors**: `factors(20) == {2, 5}` and  `factors(9) == {3}`, so 20 is valid but not 9.
# - **Product**: the result of multiplying numbers together, e.g. `prod((6, 60)) == 360`.
# 
# An implementation of the key concepts:
# 

# In[1]:


from typing import *

numbers   = range(1, 71)
Selection = Tuple[int, ...] # 5 ints in the full puzzle
Candidate = Set[Selection]
Solution  = Set[Selection]

def factors(n) -> Set[int]:
    """The set of distinct prime factors of n."""
    if n == 1:
        return set()
    else:
        p = next(i for i in range(2, n + 1) if n % i == 0)
        return {p, *factors(n // p)}
    
def prod(numbers: Iterable[int]) -> int:
    """The product of a collection of numbers."""
    result = 1
    for n in numbers:
        result *= n
    return result


# # Brute force solution?
# 
# Could we generate every possible candidate, and test each one? There are (70 choose 5)<sup>5</sup> or [about](https://www.google.com/search?q=%2870+choose+5%29^5) $10^{35}$ candidate solutions, so: **NO**.  
# 
# We'll have to be more clever. I have three ideas; I will consider only:
# 
# - Numbers with at least 2 prime factors.
# - Factors with at least 5 valid numbers.
# - Sets of 25 numbers whose product is a fifth power.

# # Numbers with at least 2 prime factors
# 
# Every valid numbers must have at least two distinct prime factors (by statement 3):

# In[2]:


numbers = [n for n in numbers if len(factors(n)) >= 2]

len(numbers)


# Good! We got it down from 70 to 41 valid numbers.
# 
# # Factors with at least 5 valid numbers
# 
# All five players make a selection with the same product (by statement 4). Therefore, if one player selects a number that has the factor $p$, then that player's product will be divisible by $p$, and therefore every player must select some number that has the factor $p$, otherwise their product would be different.
# 
# For example, the valid numbers with 11 as a factor are {22, 33, 44, 55, 66}, so if one player selects one of those numbers, then the other players must select the others.
# 
# The valid numbers with 13 as a factor are {26, 39, 52, 65}, so no player can select any of those numbers, because there aren't enough of them for all five players.
# 
# So let's count  how many valid numbers each prime factor appears in:

# In[3]:


Counter(p for n in numbers for p in factors(n))


# Only the prime factors `{2, 3, 5, 7, 11}` have a count of at least 5. Let's update the set of valid numbers to contain only numbers made from valid factors:

# In[4]:


valid_factors = {2, 3, 5, 7, 11}

numbers = [n for n in numbers if factors(n).issubset(valid_factors)]

len(numbers)


# Great! Now we're down to 28 numbers! 
# 
# # Sets of 25 numbers whose product is a fifth power
# 
# The five players will together select 25 of these 28 valid numbers, and since there are only (28 choose 25) = 3,276 combinations of 25 numbers, we can quickly check to see which ones might lead to  solutions, according to this reasoning:
# - Every player's selection of 5 numbers has the same product, which we are calling $P$.
# - Therefore the product of all 25 numbers in a solution must be $P^5$ (although we don't yet know what $P$ is).
# - A set of 25 numbers whose product is not a perfect fifth power **cannot** lead to any solutions.
# - A set of 25 numbers whose product is a perfect fifth power **might** lead to solution(s).
# 
# We can check which combinations of 25 numbers multiply to a fifth power:

# In[5]:


from itertools import combinations

def is_fifth_power(i: int) -> bool: return i == round(i ** (1/5)) ** 5

result = [c for c in combinations(numbers, 25) if is_fifth_power(prod(c))]

len(result)


# There's only one combination of 25 numbers that works! 
# 
# That's good news; we can use the `result` to update `numbers` and compute  $P$:

# In[6]:


numbers = result[0]
P = round(prod(numbers) ** (1/5))


# # Answer to the first question

# In[7]:


P


# **19,958,400** is "the unique product $P$ that each player's 5 numbers multiply to." 
# 
# At this point we know the 25 numbers  that will form any solution:

# In[8]:


print(*numbers)


# However, we haven't found any solutions yet and we don't know how many solutions there are.

# # Finding solutions
# 
# `solutions(numbers, P, r)` will search through combinations of `numbers` to generate all solutions such that every selection in the solution consists of `r` numbers with product $P$. 
# 
# `solutions` yields the empty solution if there are no numbers remaining to choose from. Otherwise it considers each way to combine a `first` selection with a set of `rest` selections, where the `first` is any `r` numbers and the `rest` is any set of selections formed from the numbers that 
# were not used in `first`, and are lexicographically greater than `first` (so that we don't generate multiple permutations of the same solution). 

# In[9]:


def solutions(numbers, P, r=5) -> Iterator[Solution]:
    """Yield solutions that are selections of `r` `numbers` with product `P`."""
    if not numbers:
        yield set()
    else:
        yield from ({first, *rest}
            for first in combinations(numbers, r)
            if prod(first) == P
            for rest in solutions([n for n in numbers 
                                   if n > first[0] and n not in first], P, r))


# We can very quickly find a solution:

# In[10]:


get_ipython().run_line_magic('time', 'next(solutions(numbers, P))')


# But it takes longer (about 6 seconds) to find all the solutions:

# In[11]:


get_ipython().run_line_magic('time', 'all_solutions = list(solutions(numbers, P))')


# # Answer to the second question

# In[12]:


len(all_solutions)


# **12,781** is "how many different ways could the selections be made so that the statements are true." 
# 
# (There is some ambiguity about what "different" means: 12,781 is the answer if you think of a set of players each choosing a set of numbers. If the ordering of the players matters, multiply by $5! = 120$, and if the ordering of the numbers in each selection matters, multiply by $5!^5$.

# # To do
# 
# - You could explore solutions with different numbers of players (default 5), selected numbers per player (default 5), range of valid numbers (default 1 to 70), and minimum number of distinct factors for selected numbers (default 2).
# 
# 
# - You could answer the trivia question "why didn't I just import [`numpy.prod`](https://numpy.org/doc/stable/reference/generated/numpy.prod.html)?"
# 
# 
# - If 6 seconds is too long to wait, you could speed up `solutions` by changing from the generate-and-test approach of looking at all `combinations(numbers, r)` and then checking to see if the product is `P`, to an incremental approach that only considers partial results that could lead to a product `P`.
# 
# 
# - You could implement a completely different approach to solving the problem (actually the one I thought of first): 
#   1. Get the valid numbers down to 28.
#   2. Consider all (28 choose 5) = 98,280 possible selections of 5 numbers.
#   3. Group selections by their products, e.g., a dict: `{19958400: [(6, 15, 56, 60, 66), ...]}`.
#   4. For each product in the dict that has at least 5 selections in its list:
#     - Find all combinations of 5 selections that consist of  distinct numbers. 
#     
# 
# - You could add more unit tests to the following meager test suite:

# In[13]:


def is_solution(candidate) -> bool:
    """Is this a valid solution?"""
    numbers  = [n for selection in candidate for n in selection]
    products = {prod(selection) for selection in candidate}
    return (len(numbers) == len(set(numbers))               # Statement 1
            and all(len(factors(n)) >= 2 for n in numbers)  # Statement 3
            and len(products) == 1)                         # Statement 4

assert factors(9) == {3}
assert factors(10) == {2, 5}
assert factors(42) == {2, 3, 7}
assert factors(41) == {41}
assert factors(1) == set()
assert factors(168000) == {2, 3, 5, 7}

assert prod([2, 3, 7]) == 42
assert prod([41]) == 41
assert prod([]) == 1
assert prod([2, 3, 2, 5, 2, 5, 2, 5, 2, 7, 2]) == 168000

assert     is_fifth_power(100000)
assert     is_fifth_power(1234567890 ** 5)
assert not is_fifth_power(99999)
assert not is_fifth_power(1234567890 ** 5 + 1)

assert P == 19958400
assert numbers == (6, 10, 12, 14, 15, 18, 20, 21, 22, 24, 28, 30, 33, 
                   36, 40, 42, 44, 45, 48, 50, 54, 55, 56, 60, 66)
assert len(all_solutions) == 12781
assert all(is_solution(s) for s in all_solutions)

assert     is_solution({(6, 60), (10, 36), (12, 30), (15, 24), (18, 20)})
assert not is_solution({(6, 60), (10, 60), (12, 30), (15, 24), (18, 20)}) #1
assert not is_solution({(9, 40), (10, 36), (12, 30), (15, 24), (18, 20)}) #3
assert not is_solution({(7, 60), (10, 36), (12, 30), (15, 24), (18, 20)}) #4