#include "prison.h"
#include <vector>
// DFA state structure.
struct DFAState {
int id; // The state's unique id in 0..(numStates-1)
int digit; // The digit index (0 = most-significant; increases by one each step)
int expected; // The expected trit from bag A (0, 1, or 2)
int turn; // Which bag to inspect: 0 means check bag A, 1 means check bag B.
int len; // The weight for this digit, i.e. p3[D-1-digit] (a power of 3)
};
// Transition function for a state s and coin count k.
// For k coming from the bag the DFA instructs you to inspect.
int transition(const DFAState &s, int k) {
// For the initial state (id == 0), we always go to state 1.
if (s.id == 0) return 1;
// Compute the current digit value (t) of the coin count.
int t = (k / s.len) % 3;
// If coin's trit is less than the expected trit, return a terminal output.
if (t < s.expected)
return (s.turn == 0 ? -1 : -2);
// If coin's trit is greater than the expected trit, return the opposite terminal output.
else if (t > s.expected)
return (s.turn == 0 ? -2 : -1);
else {
// When the coins have the same trit in this digit:
if (s.len == 1)
return 0; // Should never happen since coins are distinct.
// Compute the next state's "expected" value.
int nextExpected = (k * 3 / s.len) % 3;
// Next state corresponds to the next digit; states are arranged so that
// for digit index i (0-based) the three states are at positions (i*3 + j + 1).
int nextState = (s.digit + 1) * 3 + nextExpected + 1;
return nextState;
}
}
std::vector<std::vector<int>> devise_strategy(int N) {
// Build vector of powers of 3.
std::vector<int> p3 = {1};
while (p3.back() <= N) {
p3.push_back(p3.back() * 3);
}
// D is the number of digits needed (with p3[D] > N).
int D = p3.size() - 1;
// Total states: one initial state plus 3 states per digit.
int numStates = 3 * D + 1;
// Build the DFA state table.
std::vector<DFAState> dfa(numStates);
// State 0 is the initial state.
dfa[0] = DFAState{0, -1, -1, 0, 0}; // Dummy values; state 0 is handled separately.
// For each digit position i (i = 0 corresponds to the most significant digit)
// and each expected trit value j in {0, 1, 2}:
for (int i = 0; i < D; ++i) {
for (int j = 0; j < 3; ++j) {
int stateId = i * 3 + j + 1; // States 1..(3*D)
dfa[stateId] = DFAState{
stateId,
i, // current digit index (0 is most significant)
j, // expected trit from bag A
i % 2, // turn alternates: 0 for i even, 1 for i odd.
p3[D - 1 - i] // weight = 3^(D-1-i)
};
}
}
// Build the strategy table (the DFA transitions).
// The table has one row for each state (whiteboard number) and N+1 columns.
// Column 0 indicates which bag to inspect:
// for state 0 we use bag A (0),
// for other states we use the state's "turn" (0 for bag A; 1 for bag B).
// Columns 1..N contain either:
// a nonnegative next state's id (if not terminal),
// or -1 (to identify bag A as having fewer coins)
// or -2 (to identify bag B as having fewer coins).
std::vector<std::vector<int>> strategy(numStates, std::vector<int>(N + 1, 0));
for (int s = 0; s < numStates; ++s) {
// Determine which bag to inspect.
int bagToInspect = (s == 0 ? 0 : dfa[s].turn);
strategy[s][0] = bagToInspect;
for (int k = 1; k <= N; ++k) {
strategy[s][k] = transition(dfa[s], k);
}
}
return strategy;
}
