#include "choreography.h"
#include <vector>
#include <numeric>
// Global variables to hold the state
int N_dancers;
std::vector<int> initial_p; // initial_p[initial_pos] = dancer_id
std::vector<int> initial_pos; // initial_pos[dancer_id] = initial_pos
// The transformation state is defined by three components:
// 1. A base permutation `current_f`
// 2. A swap mask `xor_x`
// 3. A cyclic shift `shift_s`
// A dancer who started at an initial position `p` is now at the physical position:
// (current_f[p] ^ xor_x + shift_s) % N
std::vector<int> current_f;
int shift_s;
int xor_x;
// Initializes the state.
// Note: The signature is changed to match your choreography.h file.
void init(int N, std::vector<int> P) {
N_dancers = N;
initial_p.resize(N);
initial_pos.resize(N);
current_f.resize(N);
for (int i = 0; i < N; ++i) {
initial_p[i] = P[i];
initial_pos[P[i]] = i;
}
// Initially, the transformation is an identity transformation.
// F[p] = p, no shift, no swap.
std::iota(current_f.begin(), current_f.end(), 0); // Fills with 0, 1, 2,...
shift_s = 0;
xor_x = 0;
}
// Adds a move of type 1 (right shift) by lazily updating the shift.
void move_right(int K) {
// O(1) update
shift_s = (shift_s + K) % N_dancers;
}
// Adds a move of type 2 (left shift) by lazily updating the shift.
void move_left(int K) {
// O(1) update. The addition of N_dancers handles potential negative results from subtraction.
shift_s = (shift_s - K + N_dancers) % N_dancers;
}
// Adds a move of type 3 (swap adjacent) by lazily flipping the xor mask.
void swap_places() {
// O(1) update
xor_x ^= 1;
}
// Adds a move of type 4. This is the only O(N) operation.
// It "bakes" the lazy transformations (shift and swap) into the base permutation `current_f`.
void move_around() {
// Step 1: Calculate the current physical state of the board.
// We create p_current where p_current[physical_pos] = dancer_id.
std::vector<int> p_current(N_dancers);
for (int i = 0; i < N_dancers; ++i) {
int dancer_who_started_at_i = initial_p[i];
// Corrected modulo arithmetic to prevent negative results.
long long base_pos = (long long)current_f[i] ^ xor_x;
int physical_pos = (base_pos + shift_s) % N_dancers;
p_current[physical_pos] = dancer_who_started_at_i;
}
// Step 2: Determine the new permutation based on the `move_around` rule.
// The rule is: new position of dancer `d` is `p_current[d]`.
// We create `next_f` to store this new base permutation.
std::vector<int> next_f(N_dancers);
for (int d = 0; d < N_dancers; ++d) {
int p_init = initial_pos[d]; // The initial position of dancer d
int new_pos = p_current[d]; // The new physical position for dancer d
// The new logical position for the dancer who started at p_init is new_pos.
next_f[p_init] = new_pos;
}
// Step 3: Update the state.
// The new base permutation is `next_f`.
current_f = next_f;
// Reset the lazy transformations, as their effects are now baked into `current_f`.
shift_s = 0;
xor_x = 0;
}
// Returns the current position of dancer D.
int get_position(int D) {
// O(1) calculation by applying the lazy transformations to the base permutation.
int p_init = initial_pos[D];
int p_permuted = current_f[p_init];
int p_swapped = p_permuted ^ xor_x;
// Use long long to prevent overflow before the modulo operation
long long p_shifted = (long long)p_swapped + shift_s;
int p_final = (p_shifted % N_dancers + N_dancers) % N_dancers;
return p_final;
}
