Algorithm (Segment Tree)
- For $i$ with $A[i]\leq i,$ there are two possible intervals of $K$ that will result in gaining points: $[0,i-A[i]]$ and $[i+1,N-1]$.
- For $i$ with $A[i]>i,$ there is one interval available: $[i+1,N-A[i]+i].$
- So we can maintain a segment tree with max and plus monoids to keep track of points gain for possible intervals. In the end, we find the smallest index that matches the highest point gain. For details see code.
Code
template <class T>
struct plus_monoid {
typedef T value_type;
value_type unit() const { return value_type(); }
value_type mult(value_type a, value_type b) const { return a + b; }
};
template <class T>
struct max_monoid {
typedef T value_type;
value_type unit() const { return std::numeric_limits<T>::lowest(); }
value_type mult(value_type a, value_type b) const { return std::max(a, b); }
};
template <class T>
struct plus_max_action {
typename max_monoid<T>::value_type operator () (typename plus_monoid<T>::value_type f,
typename max_monoid<T>::value_type x) const {
return (x == max_monoid<T>().unit() ? x : f + x);
}
};
template <class MonoidX, class MonoidF, class Action>
struct lazy_segment_tree {
static_assert(std::is_invocable_r<typename MonoidX::value_type, Action, typename MonoidF::value_type, typename MonoidX::value_type>::value, "");
typedef typename MonoidX::value_type value_type;
typedef typename MonoidF::value_type operator_type;
static constexpr auto mon_x = MonoidX{};
static constexpr auto mon_f = MonoidF{};
static constexpr auto act = Action{};
mutable std::size_t n;
mutable std::vector<value_type> a;
mutable std::vector<operator_type> f;
lazy_segment_tree() = default;
lazy_segment_tree(size_t a_n) { resize(a_n); }
void resize(size_t a_n, std::optional<value_type> value = {}) {
n = 1 << (32 - __builtin_clz(a_n));
a.resize(2 * n - 1, mon_x.unit());
f.resize(n - 1, mon_f.unit());
if (value)
range_set(0, n, *value);
}
template <class InputIterator>
lazy_segment_tree(InputIterator first, InputIterator last) {
resize(std::distance(first, last));
std::copy(first, last, a.begin() + (n - 1));
for (int i = (n - 1) - 1; i >= 0; --i)
a[i] = mon_x.mult(a[2 * i + 1], a[2 * i + 2]);
}
void point_set(int i, value_type b) {
range_set(i, i + 1, b);
}
void range_set(int l, int r, value_type b) {
range_set(0, 0, n, l, r, b);
}
void range_set(int i, int il, int ir, int l, int r, value_type b) {
if (l <= il and ir <= r and ir - il == 1)
a[i] = b;
else if (ir <= l or r <= il) { }
else {
range_apply(2 * i + 1, il, (il + ir) / 2, 0, n, f[i]);
range_apply(2 * i + 2, (il + ir) / 2, ir, 0, n, f[i]);
f[i] = mon_f.unit();
range_set(2 * i + 1, il, (il + ir) / 2, l, r, b);
range_set(2 * i + 2, (il + ir) / 2, ir, l, r, b);
a[i] = mon_x.mult(a[2 * i + 1], a[2 * i + 2]);
}
}
void point_apply(int i, operator_type g) {
range_apply(i, i + 1, g);
}
void range_apply(int l, int r, operator_type g) {
range_apply(0, 0, n, l, r, g);
}
void range_apply(int i, int il, int ir, int l, int r, operator_type g) {
if (l <= il and ir <= r) {
a[i] = act(g, a[i]);
if (i < f.size()) f[i] = mon_f.mult(g, f[i]);
}
else if (ir <= l or r <= il) { }
else {
range_apply(2 * i + 1, il, (il + ir) / 2, 0, n, f[i]);
range_apply(2 * i + 2, (il + ir) / 2, ir, 0, n, f[i]);
f[i] = mon_f.unit();
range_apply(2 * i + 1, il, (il + ir) / 2, l, r, g);
range_apply(2 * i + 2, (il + ir) / 2, ir, l, r, g);
a[i] = mon_x.mult(a[2 * i + 1], a[2 * i + 2]);
}
}
value_type point_get(int i) const {
return range_get(i, i + 1);
}
value_type range_get(int l, int r) const {
return range_get(0, 0, n, l, r);
}
value_type range_get(int i, int il, int ir, int l, int r) const {
if (l <= il and ir <= r)
return a[i];
else if (ir <= l or il >= r)
return mon_x.unit();
else {
const auto lacc = range_get(2 * i + 1, il, (il + ir) / 2, l, r);
const auto racc = range_get(2 * i + 2, (il + ir) / 2, ir, l, r);
return act(f[i], mon_x.mult(lacc, racc));
}
}
value_type fast_range_get(int l, int r) const {
if (l == 0 and r == n) return a[0];
value_type lacc = mon_x.unit(), racc = mon_x.unit();
for (int l1 = (l += n), r1 = (r += n) - 1; l1 > 1; l /= 2, r /= 2, l1 /= 2, r1 /= 2) {
if (l < r) {
if (l % 2 == 1) lacc = mon_x.mult(lacc, a[(l++) - 1]);
if (r % 2 == 1) racc = mon_x.mult(a[(--r) - 1], racc);
}
lacc = act(f[l1 / 2 - 1], lacc);
racc = act(f[r1 / 2 - 1], racc);
}
return mon_x.mult(lacc, racc);
}
};
class Solution {
public:
int bestRotation(vector<int>& A) {
const auto N = size(A);
const auto seg_tree = [&](lazy_segment_tree<max_monoid<int>, plus_monoid<int>, plus_max_action<int>> self = {}) {
self.resize(N, 0);
for (int i = 0; i < N; ++i) {
if (A[i] <= i) {
self.range_apply(0, i - A[i] + 1, 1);
self.range_apply(i + 1, N, 1);
}
else if (A[i] > i)
self.range_apply(i + 1, N - A[i] + i + 1, 1);
}
return self;
}();
const auto solution = [&] {
const auto target = seg_tree.range_get(0, N);
const auto match = [&] {
for (int i = 0; i < N; ++i)
if (seg_tree.point_get(i) == target)
return i;
return -1;
}();
return match;
}();
return solution;
}
};