めもめも

このブログに記載の内容は個人の見解であり、必ずしも所属組織の立場、戦略、意見を代表するものではありません。

チャイティンが実装した停止確率Ωに関するコードの覚書

何の話かと言うと

数学の限界[復刻改装版] (00)

数学の限界[復刻改装版] (00)

Amazon

上記の書籍では、チャイティンの定数 Ω の説明とそれに関連した Lisp の独自実装コードが紹介されていますが、講義録という性質上、説明がやや言葉足らずで、コードの中身を読み込んで、はじめて「なるほど」と理解できた部分がいろいろあります。その「なるほど」の部分の覚書です。（筆者個人の理解ですので、理解に誤りがあるかも知れません。）

チャイティンの定数とは

別名で「プログラムの停止確率」とも呼ばれるもので、プログラムコード p のビット数を |p| として、

　\displaystyle \Omega = \sum_{p:停止}2^{-|p|}

で与えられます。なのですが、これだけでは、なぜこれが「確率」になるのかわからないですよね・・・。まずは、プログラムの実行環境を明示する必要があります。

実は、チャイティンの実装では、Mathematica で実装した独自の拡張 Lisp インタープリターが用いられており、プログラムを表すバイナリーコード（ビット列） p は次のように解釈されます。

(1) |p| が8の倍数で無い時は、正しいプログラムコードとは見なさない
(2) |p| が8の倍数の時は、8ビットごとにASCIIコードで文字に変換して、得られた文字列が「Lisp の正しい S 式 + 終端文字 "\n"」であれば、正しいプログラムコードと見なす
(3) 得られた文字列が 「Lisp の正しい S 式 + 終端文字 "\n"」でなければ、正しいプログラムコードとは見なさない

終端文字 "\n" の存在によって、正しいプログラムコード p に対して、p の後ろにさらにビットを加えたものは、決して、正しいプログラムコードにはなりません。（正しい S 式の途中には終端文字 "\n" は含まれることはありません。）この点に注意すると、この意味での「正しいプログラム」の集合を P として、

　\displaystyle \sum_{p\in P}2^{-|p|} \le 1

が成り立つことが分かります。

たとえば、「0」「1」がどちらも正しいプログラムだとすると（本当は正しくありませんが・・・）、2ビット以上の正しいプログラムは存在しないので、

　\displaystyle \sum_{p\in P}2^{-|p|} = \frac{1}{2} + \frac{1}{2} = 1

となります。

あるいは、「0」と「10」が正しいプログラムのすべて（先頭2ビットが「11」はすべて正しいプログラムでは無い）とすると、

　\displaystyle \sum_{p\in P}2^{-|p|} = \frac{1}{2} + \frac{1}{4} = \frac{3}{4}

となります。

さらに、次のような思考実験をしてみましょう。何度もコインを投げてランダムなビット列を生成していって、

・正しいプログラム p が得られる
・これ以上続けても、永遠に正しいプログラムにはならないビット列（上記の「11」のような場合）が得られる

のどちらかになったらそこでコイン投げを停止します。この方法で、正しいプログラム p が得られる確率は \displaystyle 2^{-|p|} ですので、

　\displaystyle \sum_{p\in P}2^{-|p|} \le 1

は、まさに（この思考実験のプロセスにおいて）「正しいプログラムが得られる確率」に一致します。

そしてさらに、正しいプログラム全体は、「停止するプログラム（S 式の評価が有限時間で完了する）」と「停止しないプログラム（S 式の評価が永遠に完了しない）」に分類されますので、（上記の思考実験のプロセスにおいて）「停止するプログラムが得られる確率」が冒頭で説明した「プログラムの停止確率 Ω」になります。

　\displaystyle \Omega = \sum_{p:停止}2^{-|p|}

終端文字の導入によって、「正しいプログラムにビットを付け加えたものは決して正しいプログラムにならない」という部分は上記の書籍を理解するかなり大きなポイントですが、書籍内でほとんど説明されていないので注意が必要です。（唯一あるとすれば、p.21 で受講者の質問に対して「妥当なプログラムの拡張は、妥当なプログラムにはなりません」と答えている部分でしょうか。）

Ωを計算する方法

Ωを二進数で表記すると、

\Omega = {0.00000010010010110...}_{(2)}

のような値になると想像できますが、特定の N について、「小数点以下 N 桁目までの値」を計算することは可能でしょうか？

たとえば、1ビットのプログラム「0」が（本当にはそうではありませんが仮に）正しいプログラムで、実行してみたら（運よく）停止したとします。この事実から、

\frac{1}{2} \le \Omega \le 1

であり、

\Omega = {0.1...}_{(2)}

と1桁目が決定できます。（ \Omega = {1.0}_{(2)} の可能性もありますが、その場合は、無限小数展開で \Omega = {0.1111...}_{(2)} と表記するものとします。）

そこで、プログラムのサイズを1ビットずつ大きくしながら、すべての正しいプログラムが停止するか確かめていけば、Ωの値を1桁ずつ決定していくことができるでしょうか？

残念ながらそうは行きません。

先ほどは、「0」が正しいプログラムで、さらに（運良く）停止したのでよかったのですが、もしも、「0」も「1」も正しいプログラムでなければ、さらに先のビットを見ないとΩの1桁目は決まりません。また、あるビット数で正しいプログラムが得られた際に、これを実行してもいつまでたっても停止しなかったとします。プログラムが停止するかを判定する汎用のアルゴリズムはありませんので、（プログラムの内容によっては）実行時間が足りなくてまだ停止しないだけなのか、本当に停止しないのかが判別できないことが起こり得ます。そうなると、この手続きはこれ以上先に進めなくなるので、「小数点以下 N 桁目までの値」がどうがんばっても計算できなくなる N が必ず存在することになります。

ちなみに、「これ以外に何かもっとうまい手続きがあるかも・・・」と思うかも知れませんが、それが存在しないことが、冒頭の書籍で証明されている、

　\displaystyle H(\Omega_N)>N-8000

という関係から分かることになります。（この証明はあとで紹介します。）

単調増加列としてΩの下限を与えるコード

Ωの各桁を個別に計算することはできないと言いましたが、「Ωは少なくともX以上である」というXを計算することはできます。また、計算時間を長くすれば、Xの値を下から上に（Ωの真の値に）徐々に近づけていくことができます。具体的には、引数 n を取る、次のような関数 Omega(n) が実装されています。

(1) n ビットのコードをすべて生成する
(2) 生成したコードを順番に実行する。ただし、制限時間を n として、n ステップの処理で終了しなかった場合は、そこで打ち切る。（正しいプログラムでない場合は、実行時エラーで終了する）
(3) 打ち切られずに正常終了したコード p について、\displaystyle 2^{-|p|} を足し上げた値を返す

簡単ですね。ただし、微妙な注意点があります（書籍の中では明確には説明されていませんが・・・）。先ほど、「正しいプログラムコードの条件」として、

・8ビットごとにASCIIコードで文字に変換して、得られた文字列が「Lisp の正しい S 式 + 終端文字 "\n"」であれば、正しいプログラムコードと見なす

と説明しました。チャイティンの実装したコード（Mathematica で実装した Lisp インタープリター）をよーーーーく読み解くと、文字列に変換した際に「正しい S 式 "\n" 余分なコード」となるビット列が与えられた場合、上記の (2) のステップでは、前半の「正しい S 式 "\n"」のみを取り出して実行するようになっています。つまり、(1) のステップで生成されるビット列には、実質的には、「n ビット以下のすべての正しいコード」が含まれており、Omega(n) は、

・（制限時間 n の範囲で）停止する n ビット以下のすべての正しいコード

を洗い出すようになっているのです。

そのため、n を大きくしていくと、「正常終了が確認できるコードの集合」には新たな要素が追加されていき、Omega(n) の出力も単調に増加していきます。そして、停止可能な任意のプログラムコードは、十分大きな n で必ず検出されるはずですので（この n がどのぐらい大きければよいかが分からない所が本質的な問題ではあるのですが）、理論上は、n \to \infty の極限で、Omega(n) は Ω に収束することになります。

で、実際にチャイティンが実装した Omega(n) の実行結果を見ると、

・Omega(0) 〜 Omega(7)：すべて 0
・Omega(8) = 1 / 256

という結果になっています。落ち着いて考えるとわかるのですが、7ビット以下では正しいプログラムコードは１つも表現できないので、Omega(0) 〜 Omega(7) がすべて 0 なのはその通りですね。（実行したら即座にエラーで終了する。）一方、8ビット以下の場合、1つだけ正しいプログラムコードが作れます。何でしょう・・・・

そう。「"\n"（終端文字のみ）」です。

つまり中身のまったくないコードです。ただ、Lisp のコードとしては、正しく実行することができて、実行結果は、「nil」もしくは「()」となります。これが、Omega(8) = 1 / 256（=2^{-8}）の正体です。これにより、（チャイティンが実装した Lisp の実装においては）

　\displaystyle 2^{-8} (= {0.00000001}_{(2)}) \le \Omega \le 1

が証明できたことになります。

ただし、繰り返しになりますが、Ωの小数点以下8桁目までが {0.00000001}_{(2)} と確定したわけでありません。さらに大きな n を試していけば、Omega(n) の出力値は、これらの桁が変化するほどに大きくなる可能性は残されています。（9ビット以上のコードのかなりの割合が停止しないといけないので、ほぼほぼ確定している気はしますが、証明されているわけではありません。）

Omega(n) の実装に関する注意点

割と細かい話なのですが、Omega(n) の実装をよく見ると、

・n ビットのコードをすべて生成して、（時間制限 n において）実行が正常に停止したものについて、2^{-n} を足し上げる

という処理を行っています。これをそのまま理解すると、「n ビット以下のすべての正しいコード p について、2^{-|p|} を足し上げる」という意図した処理にならないようにも思えます。

なのですが・・・これで正しい実装になっています。先ほど触れたように、チャイティンの実装したコード（Mathematica で実装した Lisp インタープリター）は、文字列に変換した際に「正しい S 式 "\n" 余分なコード」となるビット列が与えられた場合、前半の「正しい S 式 "\n"」のみを取り出して実行します。従って、「n ビット以下のすべてのコード」の中には、（p = ( A B C ) を正しいコードの1つとして）「( A B C )"\n"hoge」とか「( A B C )"\n"hoga」など、p と解釈されるコードが重複してふくまれています。これらの重複を含めて、2^{-n} を足し上げれば、結果的に、1つの p に対して 2^{-|p|} を1回加えるのと同じ結果になります。こんなトリッキーな事が文章で明示的に解説されていないのは、なかなか辛いです。。。

\Omega_n を計算するアルゴリズムが存在しない事の証明

Ωの値を小数点以下N桁目まで正確に計算するアルゴリズムはないと説明しました。

　\Omega_N = \Omega の N 桁目までの正確な値

と定義して、\Omega_N を計算するアルゴリズムは存在しないという事です。

冒頭の書籍では、（背理法に近いような）ややアクロバティックな手法でこれが証明されています。そのあらすじを紹介します。

まず、\Omega_N の値を計算するアルゴリズムが存在すると仮定して、それを実装したコードを \Omega_N^* とします。\Omega_N^* を実行（評価）すると、\Omega_N が得られます。

次に、このコードをサブルーチンとして利用して、「停止する N ビット以下のすべてのコードの実行結果」（コードそのものではなくて、実行結果である点に注意！）を網羅的に出力するコード \phi\Omega_N^* を実装します。具体的な実装方法は後ほど説明しますが、（\Omega_N^* があると仮定すれば）これは確かに実装できます。

で、さらにここで、

・プログラムコード \phi\Omega_N^* の長さが N ビット以下である ------- (*)

と仮定すると矛盾が生じます。

なぜでしょう。

「プログラムコード \phi\Omega_N^* の出力」には、「停止する N ビット以下のすべてのコードの実行結果」が含まれているので、(*) が正しければ、この中には、「プログラムコード \phi\Omega_N^* の出力」も含まれており、「集合 A の要素に集合 A 自身が含まれている」という集合論的な矛盾が生じます。まあ、実際には、「停止する N ビット以下のすべてのコードの実行結果」を生成する際に無限再帰が生じて、\phi\Omega_N^* は停止しないコードになってしまうのでしょう。

したがって、 \phi\Omega_N^* を間違いなく実装すれば、そのコードの長さは N を超えるはずです。ちなみに、チャイティンの実装では、 （サブルーチン \Omega_N^*）を除いた \phi の部分のコードは、ちょうど 8,000 ビットでした。したがって、\phi を除いた \Omega_N^* の部分は、N-8000 ビット以上であると主張することができます。この事実を情報理論の言葉で表したものが、先ほどの

　\displaystyle H(\Omega_N)>N-8000

という関係になります。\phi の長さは実装しだいで変わりますが、いずれにしても有限の値であることには変わりないので、上記の関係は、\Omega_N を計算するには、本質的に N に比例する長さのコードが必要であり、(8000 の部分が無視できるほど）十分大きな N では、実質上、\Omega_N の値をハードコーディングする以外の方法が無いことになります。つまり、事前に \Omega_N の値を知らなければ \Omega_N の値を出力するコードを書くことはできず、\Omega_N を計算するアルゴリズムは本質的に存在しないことになります。

\phi の実装

冒頭の書籍では、8,000 ビットとなる \phi の実装例が紹介されています。ざっくりと、次のような処理を実装しています。

(1) サブルーチン \Omega_N^* から \Omega_N の値を得る
(2) Omega(1), Omega(2), ... を Omega(t) \ge\Omega_N となるまで実行する
　※ この時、Omega(t) は N ビット以下の停止可能なコードをすべて発見しており、N ビット以下の停止可能なコードは、ここで得られた実行時間 t 以内に停止することが保証される（停止可能な N ビット以下のコードで Omega(t) が未発見のものがあるとすると、Omega(t) にその確率を加えた値は、\Omega_N に含まれる値を変化させるので、\Omega_N が小数点以下N桁目まで正確に計算しているという前提に矛盾する。）
(3) そこで、N ビット以下のすべてのコードを生成して、それらを順に、制限時間 t の下に実行して、停止するものをすべて発見する。
(4) 停止することが分かったコードについて、その実行結果を順番にリストに詰め込む

これがちょうど、8,000ビットのコードの中身になります。（厳密には、(4) のところで、エラーで停止したコードについてもエラーメッセージを実行結果としてリストに詰めていますが、「停止するコードの出力結果がすべて含まれている」ことが重要なのでよしとしましょう。）

で、これを実際に実行するには、サブルーチン \Omega_N^* が必要なのですが、先ほど、こんなサブルーチンは存在しないことが証明されてしまいました。書籍の中では、N=8 の計算例として（これが本当に正しいという保証はありませんが） \Omega_N={0.00000001}_{(2)} と仮定して、値「(0 0 0 0 0 0 0 1)」を直接に出力するコードを利用しています。

実行結果を見ると、8ビット以下で表現可能なすべてのコードに対して、正常終了して「()」（=nil）を出力するコードが1つだけあり、残り全ては「out of data」エラーで停止しています。これは、（チャイティンが実装したインタープリターの）「終端文字 "\n" が得られるまで1ビットずつコードを読み続ける」という実装によるもので、"\n" に対応したコードは「()」で正常終了、それ以外は、「次のビットを読もうとしてそれ以上ビットがなかった（out of data）」というパターンに対応しているものと想像できます。

ちなみに、細かい話ですが、「終端文字 "\n" が得られるまで1ビットずつコードを読み続ける」という実装から、（私のコードの読み方が間違っていなければ・・・）「A B C "\n" D」と「A B C "\n" E」は、どちらも「A B C "\n"」と解釈されるはずで、\phi は、「停止する N ビット以下のすべてのコードの実行結果」を重複なく出力するわけではなさそうです。実際に重複が発生するのは、N がかなり大きくなったところで、実行時間が指数的に増加することを考えると、実際に試すのは難しそうですが。

089 - Partitions and Inversions（★7）の解説

何の話かと言うと

atcoder.jp

上記の問題について、AVL 木（ソート済みのリストに対して、ソートを保った挿入・削除を O(\log M) で実行できるデータ構造）を用いた別解を紹介します。

再帰処理による解法

まずは直感的にわかりやすい再帰的な解法を考えます。

与えられた数列の最初の分割位置について場合分けをします。たとえば、最初の区間として可能なものが次の 3 つだとします。

・(A_1)
・(A_1,\,A_2)
・(A_1,\,A_2,\,A_3)

それぞれの場合について、

・(A_1) の場合： (A_2,\cdots,A_N) に対する解 dp[2]
・(A_1,\,A_2) の場合： (A_3,\cdots,A_N) に対する解 dp[3]
・(A_1,\,A_2,\,A_3) の場合： (A_4,\cdots,A_N) に対する解 dp[4]

があらかじめ求められているとすれば、

・dp[1] = dp[2] + dp[3] + dp[4]

が求める答えになります。

そこで一般に、dp[n] （n=1,...,N) を求める関数を solve(n) として実装すれば、再帰的に計算することができます。

最初の区間として可能な範囲は、先頭から要素を順にソート済みリストに（ソートを保ちながら）追加していき、追加する際に「後ろから飛び越える要素の数」を加えていき、これが K を超えたところで打ち切れば求まります。

bisect を利用してソート済みリストに要素を挿入していけば、挿入位置から「飛び越える要素数」を求めることができます。

この方針で実装すると下記になります。

import sys, bisect
from functools import lru_cache
sys.setrecursionlimit(10**6)
mod = 10**9 + 7

@lru_cache(maxsize=None)
def solve(n): # A[n:] に対する解を求める
  global N, K, A
  if n == N + 1:
    return 1

  result = 0
  L = []
  i = n
  skips = 0
  while i <= N:
    skips += len(L) - bisect.bisect_right(L, A[i])
    if skips > K:
      break
    bisect.insort_right(L, A[i])
    result += solve(i+1)
    result %= mod
    i += 1

  return result

def main(f):
  global N, K, A
  N, K = map(int, f.readline().split())
  A = [None] + list(map(int, f.readline().split()))
  print(solve(1))

main(sys.stdin)

計算手続きとしては問題ありませんが、このままでは TLE します。

再帰処理をやめて DP で実装する

・dp[1] = dp[2] + dp[3] + dp[4]

という関係を見ると、dp[N], dp[N-1], ... と後ろから順に求めれば、再帰処理をやめて DP で計算できることがわかります。個々の dp[n] の計算方法は先ほどと同じで、次のように実装できます。

import sys, copy, heapq, math, bisect
from collections import defaultdict, deque
mod = 10**9 + 7

def main(f):
  global N, K, A
  N, K = map(int, f.readline().split())
  A = [None] + list(map(int, f.readline().split()))

  dp = [None] * (N+2)
  dp[N+1] = 1
  for n in range(N, 0, -1):
    dp[n] = 0
    L = []
    i = n
    skips = 0
    while i <= N:
      skips += len(L) - bisect.bisect_right(L, A[i])
      if skips > K:
        break
      bisect.insort_right(L, A[i])
      dp[n] += dp[i+1]
      dp[n] %= mod
      i += 1

  print(dp[1])

main(sys.stdin)

これもまた正しい実装ですが、やはりまだ TLE します。

ここで計算量を確認してみましょう。まず、dp[N], dp[N-1], ..., dp[1] と計算していく O(N) のループが一番外側にあります。

次に、個々の dp[n] を求める際に、最初の区間として可能なものを調べていきますが、これは最悪ケースで O(n) のループとなるので、この段階で、全体として O(N^2) になります。これは確かに TLE します。

「最初の区間として可能な範囲」の計算を分離する

そこで、dp[n] の計算と、「最初の区間として可能な範囲」の計算を分離します。

仮に、n = 1,2,..., N について、

・A_n から始めた場合に、最初の区間として可能な範囲は、(A_n,\cdots,A_m) である

という情報を十分高速に決定できたとします。上記の m は、リスト skip_to に保存されているものとします。（m = skip_to[n]）

この場合、DP による dp[n] の計算は次のように簡単化されます。

  dp = [None] * (N+2)
  dp[N+1] = 1
  for n in range(N, 0, -1):
    dp[n] = 0
    for i in range(n+1, skip_to[n]+2):
      dp[n] += dp[i]

  print(dp[1])

ただし、これは、skip_to[n] の計算を外だししただけで、O(N^2) のループであることには変わりありません。

しかしながら・・・、上記は「部分和」の計算になっているので、総和の差分として求めれば、計算量を O(N) に減らすことができます。

  dp = [None] * (N+2)
  dp_sum = [None] * (N+2)
  dp[N+1] = 1
  dp_sum[N+1] = 1
  for n in range(N, 0, -1):
    if skip_to[n] + 2 > N+1:
      dp[n] = dp_sum[n+1]
    else:
      dp[n] = dp_sum[n+1] - dp_sum[skip_to[n]+2]
    dp[n] %= mod
    dp_sum[n] = dp_sum[n+1] + dp[n]
    dp_sum[n] %= mod

  print(dp[1])

というわけで、あとは、「最初の区間として可能な範囲」の計算をうまく実装できればOKです。

・先頭から要素を順にソート済みリストに（ソートを保ちながら）追加していき、追加する際に「後ろから飛び越える要素の数」を加えていき、これが K を超えたところで打ち切る

という前述の作戦を n = 1, 2, ..., N について実行するのはどうでしょうか？

ダメですね。

n = 1 に対して、お尻の部分を i = 2, 3, ... と伸ばしていって打ち切る。次に n = 2 に対して、再び、お尻の部分を i = 3, 4, ... と伸ばしていって・・・とすると、O(N^2) になってしまいます。

実は、このやり方の場合、n = 1 についてお尻を伸ばし切って打ち切った際に、再び、お尻を i = 2 にリセットする部分に無駄があります。A_1 から始まる場合のお尻と、A_2 から始まる場合のお尻を比べれば、A_2 から始める場合の方が、お尻の位置はより後ろに伸びるはずなので、毎回、お尻をリセットするのではなく、そのままお尻を伸ばし続ければよいのです。

ただし、A_1 を取り除くことによって、「飛び越えた数」をその分だけ減らす必要があります。具体的には、これまでに追加した要素で、A_1 より小さいものの個数を減らせばOKです。この個数も bisect で求めることにすれば、次のような実装が可能です。

  skip_to = [None] * (N+1)
  skips = 0
  L = [A[1]]
  i = 2
  for n in range(1, N+1):
    while True:
      if i == N+1:
        skip_to[n] = N
        break
 
      if len(L) == 0:
        L.append(A[i])
        i += 1
        continue
 
      delta = len(L) - bisect.bisect_right(L, A[i])
      if skips + delta > K:
        skip_to[n] = i - 1
        count = bisect.bisect_left(L, A[n])
        skips -= count
        del L[count]
        break
      skips += delta
      bisect.insort_right(L, A[i])
      i += 1

ここまでの実装を提出すると・・・

atcoder.jp

惜しいです・・・。1 つだけ TLE が残ります。。。。

bisect の効率化

skip_to を求める上記の実装は、表面的には、N 回のループで O(N) に見えますが、それぞれのループの中で、bisect を用いた処理を行なっています。bisect の処理は O(N\log N) なので、全体として O(N^2\log N) なんですね。残念。

しかしながら、bisect と同等の処理を O(\log N) で実施する AVL 木がありました。

enakai00.hatenablog.com

bisect を AVL 木に置き換えれば、無事に AC です。

atcoder.jp

AVL木の実装は、個人の方がブログで公開しているこちらの実装を使っています。

・PythonでAVL木を競プロ用に実装した

この実装は同一の値を保存できないので、（上記の AC するコードでは）A_1, A_2, \cdots に少しずつ小さな値を加える工夫をしてあります。

081 - Friendly Group（★5）の解説

何の話かというと

atcoder.jp

上記の問題の別解です。

公式解説 では、身長と体重を 2 次元にプロットして、身長と体重を対称に扱う解法が示されています。あえてこれらを非対称に扱うことも可能です。

身長だけの場合

仮に、身長についてだけ考えればよいとすれば、割と簡単な問題です。事前にソートした上で、いわゆる「尺取り法」で実装できます。

  A.sort()
  head = tail = 0
  a_max = a_min = A[head]

  max_len = 0
  while True:
    while a_max - a_min > K:
      tail += 1
      a_min = A[tail]
    while a_max - a_min <= K:
      max_len = max(max_len, head - tail + 1)
      head += 1
      if head == N:
        print(max_len)
        return
      a_max = A[head]

体重を含めた場合

上記の尺取り法の中では、「身長の観点で許容されるグループ」が自然に網羅されています。許容されるグループの中で、人数が最大になるものを検索する形になります。

そこで、「身長の観点で許容されるグループ」のそれぞれに対して、そのメンバーから「体重の観点でも許容されるサブグループ」を構成するという方法が考えられます。

極端な方法としては、「身長の観点で許容されるグループ」のそれぞれに対して、体重についての尺取り法を実行するというやり方が考えられます。

atcoder.jp

上記の実装では、tail を固定した状態で、head を伸ばしていき、最大限に伸ばし切ったタイミングで、体重についての尺取り法を呼び出すという実装になっています。いくつか TLE していますが、方向性としては悪くなさそうです。

TLE の 1 つの要因として、体重についての尺取り法を呼び出す際に、対象となるサブグループのコピーが発生する点がちょっと重い気がします。

そこで、体重に関しては、配列 num_B[i] を用意して、今考えている「身長の観点で許容されるグループ」の中で、

・num_B[i] = 体重が i 〜 i + K の範囲の人数

をトラッキングする様にします。こうすれば、max(num_B) として、身長・体重の両方について許容される最大人数を得ることができます。max(B) を計算するタイミングを最小限になるようにうまく調整すると無事に AC します。

atcoder.jp