//Solution by: LLM
#include <bits/stdc++.h>
#include "Scales.h"
using namespace std;
#define FOR(i, l, r) for(int i = l; i <= r; i++)
#define pb push_back
// Store all 720 possible permutations of [1, 2, 3, 4, 5, 6]
vector<vector<int>> perms;
// Define a query structure
struct Query {
char type; // 'L' (Lightest), 'H' (Heaviest), 'M' (Median), 'N' (Next Lightest)
int a, b, c, d; // Coins involved
};
vector<Query> all_queries;
// precomp[q][p] tells us the result (0, 1, or 2) of query q on permutation p
// We map the actual coin returned to branch 0, 1, or 2 to split the state space.
vector<vector<int>> precomp;
map<vector<int>, string> memo;
// Simulate the result of a query on a specific permutation
// A permutation here is the order of coins from lightest to heaviest
int simulate(const Query& q, const vector<int>& p) {
// Find the positions (weights) of the queried coins
int wa = find(p.begin(), p.end(), q.a) - p.begin();
int wb = find(p.begin(), p.end(), q.b) - p.begin();
int wc = find(p.begin(), p.end(), q.c) - p.begin();
vector<pair<int, int>> w = {{wa, q.a}, {wb, q.b}, {wc, q.c}};
sort(w.begin(), w.end()); // Sorted by weight
int ans_coin = 0;
if (q.type == 'L') ans_coin = w[0].second;
else if (q.type == 'H') ans_coin = w[2].second;
else if (q.type == 'M') ans_coin = w[1].second;
else if (q.type == 'N') {
int wd = find(p.begin(), p.end(), q.d) - p.begin();
ans_coin = -1;
FOR(i, 0, 2) {
if (w[i].first > wd) {
ans_coin = w[i].second;
break;
}
}
if (ans_coin == -1) ans_coin = w[0].second; // Wraps around
}
// Map the returned coin to branch 0, 1, or 2 based on the input order
if (ans_coin == q.a) return 0;
if (ans_coin == q.b) return 1;
return 2;
}
// Precompute all valid queries and their results on all permutations
void init_queries() {
FOR(a, 1, 6) FOR(b, a + 1, 6) FOR(c, b + 1, 6) {
all_queries.pb({'L', a, b, c, 0});
all_queries.pb({'H', a, b, c, 0});
all_queries.pb({'M', a, b, c, 0});
FOR(d, 1, 6) {
if (d != a && d != b && d != c) {
all_queries.pb({'N', a, b, c, d});
}
}
}
precomp.assign(all_queries.size(), vector<int>(perms.size()));
FOR(i, 0, (int)all_queries.size() - 1) {
FOR(j, 0, (int)perms.size() - 1) {
precomp[i][j] = simulate(all_queries[i], perms[j]);
}
}
}
// The core DFS generator
string solve(vector<int> cands, int depth) {
if (cands.size() == 1) {
// Base case: we isolated exactly one permutation
vector<int> res = perms[cands[0]];
string code = "int W[] = {" + to_string(res[0]) + ", " + to_string(res[1]) + ", " +
to_string(res[2]) + ", " + to_string(res[3]) + ", " +
to_string(res[4]) + ", " + to_string(res[5]) + "};\nanswer(W);\n";
return code;
}
// Math pruning rule: if we have more candidates than 3^(remaining queries), abort
int p3 = 1;
FOR(i, 1, 6 - depth) p3 *= 3;
if (cands.size() > p3) return "";
if (memo.count(cands)) return memo[cands];
// Try all queries to find one that splits the candidates validly
FOR(i, 0, (int)all_queries.size() - 1) {
vector<int> buckets[3];
for (int idx : cands) {
buckets[precomp[i][idx]].pb(idx);
}
// Minor optimization: don't waste a query if it doesn't split the state at all
if (buckets[0].size() == cands.size() ||
buckets[1].size() == cands.size() ||
buckets[2].size() == cands.size()) continue;
// Recursively solve for each bucket
string code0 = solve(buckets[0], depth + 1);
if (code0 == "" && buckets[0].size() > 0) continue;
string code1 = solve(buckets[1], depth + 1);
if (code1 == "" && buckets[1].size() > 0) continue;
string code2 = solve(buckets[2], depth + 1);
if (code2 == "" && buckets[2].size() > 0) continue;
// If we found a valid split where all branches are solvable, generate the C++ block
Query q = all_queries[i];
string func = (q.type == 'L') ? "getLightest" : (q.type == 'H') ? "getHeaviest" :
(q.type == 'M') ? "getMedian" : "getNextLightest";
string args = to_string(q.a) + ", " + to_string(q.b) + ", " + to_string(q.c);
if (q.type == 'N') args += ", " + to_string(q.d);
string indent(depth * 4, ' ');
string res = "int res" + to_string(depth) + " = " + func + "(" + args + ");\n";
// Branch 0 (returned q.a)
if (buckets[0].size() > 0) {
res += indent + "if (res" + to_string(depth) + " == " + to_string(q.a) + ") {\n";
res += indent + " " + code0;
res += indent + "}\n";
}
// Branch 1 (returned q.b)
if (buckets[1].size() > 0) {
res += indent + (buckets[0].size() > 0 ? "else " : "") + "if (res" + to_string(depth) + " == " + to_string(q.b) + ") {\n";
res += indent + " " + code1;
res += indent + "}\n";
}
// Branch 2 (returned q.c)
if (buckets[2].size() > 0) {
res += indent + ((buckets[0].size() > 0 || buckets[1].size() > 0) ? "else " : "") + "{\n";
res += indent + " " + code2;
res += indent + "}\n";
}
return memo[cands] = res;
}
return memo[cands] = ""; // Unsolvable branch
}
