#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 group_number = 0, N_ = N;
N -= 1;
int nearest_power_of_two = 1;
while(N != 0)
{
nearest_power_of_two *= 2;
N /= 2;
group_number += 1;
}
int x = group_number * 3;
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 == 0) // to check box A
{
ans[i][0] = 0;
}
else // to check box B
{
ans[i][0] = 1;
}
int cur_group = i / 3;
for(int j = 1; j <= N; j++)
{
int cur_N = nearest_power_of_two, group = 0;
while(group != cur_group)
{
group += 1;
cur_N /= 2;
}
int new_j = j % cur_N;
if(new_j == 0)
{
new_j = cur_N;
}
if(new_j <= cur_N / 2) // part 0
{
if(i % 3 != 0) // checked box B
{
// compare and put result for box A
if(i % 3 == 2) // A was in part 1
{
ans[i][j] = -2;
}
else // A was in part 0, move to next group
{
ans[i][j] = (cur_group + 1) * 3;
}
}
else // checked box A
{
ans[i][j] = cur_group * 3 + 1; // to check box B, box A was in part 0
}
}
else // part 1
{
if(i % 3 != 0) // checked box B
{
if(i % 3 == 1) // box A was in part 0
{
ans[i][j] = -1;
}
else // A was in part 1, move to next group
{
ans[i][j] = (cur_group + 1) * 3;
}
}
else // to check box B, box A was in part 1
{
ans[i][j] = cur_group * 3 + 2;
}
}
}
for(int j = 1; j <= N; j++)
{
if(ans[i][j] >= x)
{
ans[i][j] -= 3;
}
}
}
return ans;
}
