Subset-based solution to the 3-prisoner number sum problem

prisonersubset.cpp

#include <algorithm>
#include <iostream>
#include <numeric>
#include <set>
#include <vector>

typedef std::vector<int> triplet;

// List all ways to choose K-element subsets of an N-element vector
// By abusing next_permutation
template <typename T>
std::vector<std::vector<T>> choose(std::vector<T> values, int K) {
  std::vector<std::vector<T>> ret;

  int N = values.size();
  if (N < K)
    return ret;

  std::vector<int> mask(N, 0);
  std::fill(mask.rbegin(), mask.rbegin() + K, 1);

  // For all masks from 000..111 to 111..000, select the 1-elements
  do {
    std::vector<T> cur;
    for (int i = 0; i < N; ++i) {
      if (mask[i])
        cur.push_back(values[i]);
    }
    ret.push_back(cur);
  } while (std::next_permutation(mask.begin(), mask.end()));

  return ret;
}

// List all subsets of an N-element vector from all the K-elements subsets
template <typename T>
std::vector<std::vector<T>> subsets(std::vector<T> values) {
  std::vector<std::vector<T>> ret;
  for (int i = 0; i <= values.size(); ++i) {
    auto cur = choose(values, i);
    ret.insert(ret.end(), cur.begin(), cur.end());
  }
  return ret;
}

// Naive sqrt(N) primality testing
bool isprime(int N) {
  if (N < 2)
    return false;
  for (int i = 2; i * i <= N; ++i) {
    if (N % i == 0)
      return false;
  }
  return true;
}

bool isgood(const std::vector<triplet> &subset,
            const std::vector<triplet> &prime,
            const std::vector<triplet> &myposs) {

  // Only those elements in myposs can be actual answers
  // Just that the other two people do not know it
  for (auto trip : myposs) {
    // If trip is the actual answer, does there exist at least one person whose
    // sum for every triplet satisfying the subset question answer is trip sum?

    // Select the set (subset/complement) which has trip in it
    const std::vector<triplet> &subsetorprime =
        (std::find(subset.begin(), subset.end(), trip) == subset.end())
            ? prime
            : subset;

    // Can someone find the sum?
    bool flag = false;
    for (auto value : trip) {
      // If a triplet in the sum/complement set contains the value,
      // add its sum to a set
      std::set<int> sumset;
      for (auto trip2 : subsetorprime) {
        if (std::find(trip2.begin(), trip2.end(), value) != trip2.end())
          sumset.insert(std::accumulate(trip2.begin(), trip2.end(), 0));
      }

      // All triplets containing value have the same sum, they can hence find it
      if (sumset.size() == 1) {
        flag = true;
        break;
      }
    }

    // No one can find the sum if trip is the actual triplet, so subset not good
    if (!flag)
      return false;
  }

  // Subset works for all valid trips
  return true;
}

int main() {
  std::vector<int> values(9);
  std::iota(values.begin(), values.end(), 1);

  // Find all triplets in {1..9}
  std::vector<triplet> poss = choose(values, 3);

  // Erase triplets with even sum
  poss.erase(
      std::remove_if(poss.begin(), poss.end(),
                     [](const triplet &t) {
                       return std::accumulate(t.begin(), t.end(), 0) % 2 == 0;
                     }),
      poss.end());

  // Erase triplets with prime sum
  poss.erase(std::remove_if(poss.begin(), poss.end(),
                            [](const triplet &t) {
                              return isprime(
                                  std::accumulate(t.begin(), t.end(), 0));
                            }),
             poss.end());

  // Sort triplets: Important for set complement to work!
  sort(poss.begin(), poss.end());

  // I know that my number is 5, so I can reduce triplet possibilities
  // Important: Make sure that program logic does not need others to know this!
  auto myposs = poss;
  myposs.erase(std::remove_if(myposs.begin(), myposs.end(),
                              [](const triplet &t) {
                                return std::find(t.begin(), t.end(), 5) ==
                                       t.end();
                              }),
               myposs.end());

  // Using poss instead of myposs here
  // What if the additional yes/no information actually ends up helping others
  // even if you don't need it?
  auto questions = subsets(poss);
  for (auto subset : questions) {

    // Complement of subset
    std::vector<triplet> prime;
    std::set_difference(poss.begin(), poss.end(), subset.begin(), subset.end(),
                        std::inserter(prime, prime.begin()));

    if (isgood(subset, prime, myposs)) {
      for (auto trip : subset) {
        for (auto v : trip) {
          std::cout << v << "";
        }
        std::cout << " ";
      }
      std::cout << ": ";
      for (auto trip : prime) {
        for (auto v : trip) {
          std::cout << v << "";
        }
        std::cout << " ";
      }
      std::cout << "\n";
    }
  }
  return 0;
}