そこまで実装が重たいものは無かった。久々に全完できたのは嬉しい。
5乗なんてしたら、$A,B$ が幾らもいかない内にあっという間に超巨大な数になる。
実際計算してみると、連続する2数の差では $(A,B)=(120,119)$ で $A^5-B^5=1019663401$ と $10^9$ を越える。
つまり、$A$ が120以上になっても、引くことで $10^9$ 以下である $X$ を作れるような $B$ が存在しないので、考慮するのは119以下の数のみでよい。
$120^2$ なら全探索が可能である。
x = int(input())
facts = [i ** 5 for i in range(120)]
for a, a5 in enumerate(facts):
for b, b5 in enumerate(facts):
if a5 + b5 == x:
print(a, -b)
break
if a5 - b5 == x:
print(a, b)
break
else:
continue
break
仮に $X$ の制約がもっと大きくて、$A$ の取り得る範囲が $10^5$ 位になる場合は、$A$ が増加するごとに $B$ の取り得る値も増加するので、尺取り法を用いると $O(N)$ となる。
所持金とかでもいいのに、何故問題設定が身長なのか（ある人がある人の数倍の身長があるというのも変だし、そもそも身長を足すという行為が奇妙だし）
まず、番号 $1$ の人を考えてみる。
番号 $2,3,4,...$ のそれぞれとの身長の和として取るべき値は $1,2,3,...$ となる。ここから $A_1$ を引くと、番号1とペアが成立する相方 $i$ が持つべき $A_i$ の値が分かる。
番号 1 2 3 4 5 6 A 2 3 3 1 3 1 ----1を中心とする---- 和として取りたい値 - 1 2 3 4 5 ペアが成立する値 - -1 0 [1] 2 3 番号1は番号4の人とペアになれる
このように、1ずつ増加する値があって、それと等しい値は無いかを探すとき、そのままでは「位置と値」を両方考慮する必要があって大変。
そこで、最初に $A$ から1ずつ増加する値を引いておけば、ある決まった「値」さえ考慮すればよいという問題に置き換えられる。
番号 1 2 3 4 5 6 A 2 3 3 1 3 1 A' 2 2 1 -2 0 -4 ※Aから0,1,2,... を引く ----1を中心とする---- 和として取りたい値 - 0 0 0 0 0 ペアが成立する値 - -2 -2 [-2] -2 -2
番号2の人は、このままでは $A$ を $A'$ にする過程で1ずつ引きすぎているので、和として取りたい値でその分を埋め合わせする。
番号 1 2 3 4 5 6 A 2 3 3 1 3 1 A' 2 2 1 -2 0 -4 ----2を中心とする---- 和として取りたい値 - - -1 -1 -1 -1 ※1からの番号の差分だけマイナスになる ペアが成立する値 - - -4 -4 -4 [-4] 番号6の人とペアになれる
また、中心としている番号から左側は考慮しなくてよい。（成立するとしたら、より左側の数字を中心としたときに既に数えられている）
よって、右側から埋めていけばよい。
番号 1 2 3 4 5 6 A 2 3 3 1 3 1 A' 2 2 1 -2 0 -4 ----6を中心とする---- ペアが成立する値: （1からの番号の差分: -5）- (A6: 1) = -6 ans += count[-6] 自分より右側で A'i = -6 となったものの個数を答えに加算（最初なので当然0） count[-4] += 1 自分の A'i をcountに登録 ----5を中心とする---- ペアが成立する値: （1からの番号の差分: -4）- (A5: 3) = -7 ans += count[-7] 自分より右側で A'i = -7 となったものの個数を答えに加算 count[0] += 1 自分の A'i をcountに登録 ----4を中心とする---- ペアが成立する値: （1からの番号の差分: -3）- (A4: 1) = -4 ans += count[-4] 自分より右側で A'i = -4 となったものの個数を答えに加算（6とのペアが見つかる） count[-2] += 1 自分の A'i をcountに登録 ...
from collections import defaultdict
def solve(n, aaa):
fff = aaa.copy()
for i in range(n):
fff[i] -= i
ans = 0
count = defaultdict(int)
for i in range(n - 1, -1, -1):
ans += count[-i - aaa[i]]
count[fff[i]] += 1
return ans
n = int(input())
aaa = list(map(int, input().split()))
print(solve(n, aaa))
'AB' の時、$A,B$ のどちらかを1増やし、どちらかを1減らす'BC' の時、$B,C$ のどちらかを1増やし、どちらかを1減らす'AC' の時、$A,C$ のどちらかを1増やし、どちらかを1減らすざっと考えて、
これでまずいことはあるか？を考えると、以下のようなケースで、注意しないと、本来可能なのに不可能と判断してしまうことになる。
s A B C
1 1 0
AB どっちも等しいので、とりあえずAを増やす
2 0 0
BC 操作不可能！
s A B C
1 1 0
AB Bを増やしておけばよかった
0 2 0
BC
0 1 1
このケースでは、操作しようとしている変数が両方とも1の場合、次の $s_{i+1}$ に含まれる方を優先的に増やせば回避できる。
その他のケースを考えて特に見つからなかったので、ここまでの条件で貪欲にシミュレーションを実装すると通る。
一応、証明っぽいものを考えると、不可能となるのは以下に限られる。
$A+B+C = 1$ の時は選択権がないので、指定されたらそれまで。
$A+B+C \ge 2$ とすると、2番目の条件は必ず回避できる。
なぜなら、たとえば'AB'が指定されているとして、この操作直後に2箇所が0になりうるのは、以下の2つを共に満たす場合に限られる。
すると、大きい方を1減らすというルール上、$A \le B$ である必要があり、$A+B+C \ge 2$ なので、$A=1$。 つまり不可能となり得るのは $(A,B,C)=(1,1,0)$ の場合に限られる。
そして、それは「$(A,B)=(1,1)$ の場合は次に出てくる方を増やす」ようにしておけば回避可能。 唯一不可能となり得るケースも回避方法があるので、必ず回避可能である。
次の指定が今と同じ'AB'であっても、$(1,1)→(2,0)→(1,1)$ に戻るだけで、その時にどうするかはその時に考えればよいので、2手先以降は考慮する必要は無い。
n, a, b, c = map(int, input().split())
d = {'AB': (0, 1), 'AC': (0, 2), 'BC': (1, 2)}
queries = [d[input()] for _ in range(n)]
history = []
abc = [a, b, c]
for i in range(n):
x, y = queries[i]
if abc[x] == 0 and abc[y] == 0:
print('No')
exit()
if abc[x] == 1 and abc[y] == 1:
if i == n - 1 or x in queries[i + 1]:
history.append('ABC'[x])
abc[x] += 1
abc[y] -= 1
else:
history.append('ABC'[y])
abc[x] -= 1
abc[y] += 1
else:
if abc[x] > abc[y]:
history.append('ABC'[y])
abc[x] -= 1
abc[y] += 1
else:
history.append('ABC'[x])
abc[x] += 1
abc[y] -= 1
print('Yes')
print('\n'.join(history))