main.py

from functools import lru_cache
from collections import defaultdict
from math import log
from tqdm import tqdm


all_words = list(map(str.strip, open('words')))
guesses = list(map(str.strip, open('dict'))) + all_words


@lru_cache(None)
def get_reaction(word, guess):
    # print(word, guess)
    word = all_words[word]
    guess = guesses[guess]
    reaction = 0
    for i in range(5):
        if word[i] == guess[i]:
            reaction |= 1 << (i * 2)
        elif guess[i] in word and any(word[j] == guess[i] and word[j] != guess[j] for j in range(5)):
            reaction |= 1 << (i * 2 + 1)
    return reaction


def split_set(words, guess):
    splits = defaultdict(list)
    for w in words:
        splits[get_reaction(w, guess)].append(w)
    return list(splits.values())


def get_score(split):
    s = sum(map(len, split))
    return sum(len(p) / s * log(len(p) / s) for p in split)
    # return max(map(len, split))


def get_best_guess(words):
    best_score = 9e99
    best_guess = None
    for g in tqdm(range(len(guesses)), disable=len(words) < 500):
        split = split_set(words, g)
        score = get_score(split)
        if best_score > score:
            best_score = score
            best_guess = g
    # print(best_score, best_guess)
    return best_guess


def get_guess_times(words):  # average, max
    if len(words) <= 1:
        return 1, 1
    guess = get_best_guess(words)
    avg = 1.0
    m = 1.0
    splits = split_set(words, guess)
    if len(words) == len(all_words):
        print(guesses[guess], list(map(len, splits)))
    for split in splits:
        assert len(split) < len(words)
        cur_a, cur_m = get_guess_times(split)
        avg += cur_a * len(split) / len(words)
        m = max(m, cur_m + 1)
    return avg, m


for w in tqdm(range(len(all_words))):
    for g in range(len(guesses)):
        get_reaction(w, g)

print(get_guess_times(list(range(len(all_words)))))