#include "prison.h"
#include <vector>
#include <bits/stdc++.h>
using namespace std;
vector<vector<int> > devise_strategy(int N)
{
/*
[0, 1, 2, 3] means group 0 (1st check)
[4, 5, 6, 7] means group 1 (2nd check)
...
even number means to check Box A
odd number means to check Box B
*/
int N_ = N;
N -= 1;
int groups = 1;
int nearest_power_of_three = 1;
while(N != 0)
{
nearest_power_of_three *= 3;
N /= 3;
groups += 3;
}
nearest_power_of_three /= 3;
// cout << "nearest_groups_of_three = " << nearest_power_of_three << "\n";
int x = groups;
// cout << "x = " << x << "\n";
N = N_;
vector<vector<int> > ans(x, vector<int> (N + 1, -1));
// ans[i][0] = what to do if i is the number you see
// ans[i][j] st j >= 1 = what to do if i is the number
// you see and j is the number of balls that are there in the bag
for(int i = 0; i < x; i++)
{
if(i % 3 == 1)
{
nearest_power_of_three /= 3;
}
// cout << "nearest_power_of_three = " << nearest_power_of_three << "\n";
if(i == 0)
{
// always check box A
ans[i][0] = 0;
for(int j = 1; j <= N; j++)
{
ans[i][j] = (j - 1) / nearest_power_of_three + 1;
}
}
else
{
int cur_group = (i - 1) / 3;
if(cur_group % 3 == 1)
{
// check box A
ans[i][0] = 0;
}
else
{
// check box B
ans[i][0] = 1;
}
int cur_power = (nearest_power_of_three == 0) ? 1 : nearest_power_of_three * 3;
int x = (i - 1) % 3;
for(int j = 1; j <= N; j++)
{
if((j - 1) / cur_power < x)
{
ans[i][j] = (ans[i][0] == 1) ? -2 : -1;
}
else if((j - 1) / cur_power > x)
{
ans[i][j] = (ans[i][0] == 1) ? -1 : -2;
}
else
{
int holy = (j - 1) / cur_power;
holy %= nearest_power_of_three;
ans[i][j] = (cur_group * 3) + 4 + holy;
}
}
}
for(int j = 1; j <= N; j++)
{
if(ans[i][j] >= x)
{
ans[i][j] -= 3;
}
}
}
return ans;
}
