BST で k 番目に小さい値

1 番小さい値が欲しいだけなのに、全部を並べている

k 番目に小さい値を返す。そう言われて最初に浮かぶのは「全部取り出して並べ、k - 1 番目を見る」でしょう。答えは合います。ただ、ノードが 100 万個ある木から 1 番目が欲しいときも、100 万個を配列に詰めてから先頭を見ることになります。

欲しいのは 1 個です。100 万個ぶんの置き場所を用意する必要はありません。

数えながら進んで、k 個目で止める

やることは、順番に見ながら数え、k 個目に当たったらそこで打ち切るだけです。木から離れて、ただの配列で形を見ておきます。

数える変数を 1 つ外に置き、当たったらその場で返す。この形をそのまま木へ移します。違うのは「順番に見る」の中身だけです。

ここで気をつけるのは、数える変数の置き場所です。走査の中で作ると、呼ばれるたびに 0 へ戻ってしまいます。走査の外側に 1 つだけ置いて、どの呼び出しからも同じ数を触るようにしてください。Python で外側の変数を書き換えるときは nonlocal の宣言が要ります。

小さい順に見て回る道は決まっている

木のノードを小さい順に取り出す走査は、第 2 章で書いたものがそのまま使えます。左をすべて見てから自分、それから右です。次の木なら 3, 5, 7, 10, 15, 20 の順で出てきます。

この順番が、そのまま小さいほうからの順位になります。k = 1 なら 3、k = 4 なら 10、k = 6 なら 20 です。数を増やす場所は「左から戻ってきて、自分の値を見たとき」の 1 か所だけです。降りる前に数えたり、右へ進んだ後に数えたりすると、順位がまるごとずれます。

k = 1 だけが欲しいなら、左端に着いた時点で終わります。100 万個の木でも、降りた回数ぶんしか触りません。

止めたつもりで、止まっていない

打ち切りでつまずくのは、答えが決まった後も走査が続いてしまう形です。再帰の中で答えを見つけて戻っても、呼び出し元はその先の右側へ進んでいきます。戻った先が「もう終わったこと」を知らないからです。答えが決まったかどうかを表す印を 1 つ持ち、次へ進む前にそれを見てください。

順位の数え方もずれやすいところです。k = 1 が最小値で、0 番目はありません。k が 0 以下のときと、ノードの数より大きいときは -1 を返す約束です。木を最後まで見終わっても数が届かなかった、という状態をきちんと返り値に落としてください。

解説

k 番目に大きい値を聞かれたときは、右から先に見て同じように数えるだけです。作り直しは要りません。

1. in-order 走査が昇順になる BST の性質を利用する
2. k は 1 始まりの順位。1 番目 = 最小値
3. k が範囲外、または空木の場合は -1 を返す