C++で平衡二分探索木（BST）に和が0になるトリプレットが存在するか判定する方法

平衡二分探索木（BST）が与えられたとき、その木の中に合計が0になる3つのノード（トリプレット）が存在すれば true を返し、存在しなければ false を返す関数 is_valid_triplet() を作成することを考えます。この問題は、以下の制約のもとで設計する必要があります。

• 期待される時間計算量は O(n²)
• 追加で使用できる空間計算量は O(log n)

例えば、次のようなBSTが入力として与えられたとします。

この場合、出力は True になります。なぜなら、[-15, 7, 8] というトリプレットの合計が 0 になるからです。

解法のアプローチ

この問題を効率的に解く鍵は、BSTを中順走査（in-order traversal）によってソート済みの双方向連結リストに変換することです。こうすることで、両端から動かせる2つのポインタ（two-pointer technique）が使えるようになり、各ノードに対して残りの要素から「その値の符号を反転した値」と等しくなるペアを O(n) で探索できます。全体の計算量は O(n × n) = O(n²) となり、また平衡木における再帰の深さは O(log n) なので、追加空間の制約も満たします。

具体的には、以下の手順に従います。

ステップ1：BSTを双方向連結リストに変換する

関数 bst_to_doubli_list() を定義します。引数として root、head、tail を受け取ります。

• root が NULL の場合は何もせずに return します。
• root の左の子が NULL でない場合、左部分木に対して bst_to_doubli_list() を再帰的に呼び出します。
• root->left := tail と設定します。
• tail が NULL でなければ、tail->right := root と設定します。
• それ以外（tail が NULL の場合）は、head := root と設定します。
• tail := root と更新します。
• root の右の子が NULL でない場合、右部分木に対して bst_to_doubli_list() を再帰的に呼び出します。

ステップ2：双方向連結リスト内で2つのポインタによりペアを探索する

関数 is_in_double_list() を定義します。引数として head、tail、sum を受け取ります。

• head ≠ tail の間、以下を繰り返します。
• current := head のキー + tail のキー を計算します。
• current == sum であれば、true を返します。
• current > sum であれば、tail := tail の left ポインタ と進めます。
• それ以外の場合は、head := head の right ポインタ と進めます。
• ループが終了したら false を返します。

ステップ3：メイン処理での組み合わせ

• root が NULL であれば、false を返します。
• head = NULL、tail = NULL で初期化します。
• bst_to_doubli_list(root, &head, &tail) を呼び出し、BSTをソート済みの双方向連結リストへ変換します。
• (head の right が tail ではない) かつ (head のキー < 0) の間、以下を繰り返します。
• is_in_double_list(head の right, tail, head のキー × (-1)) が true を返せば、true を返します。
• そうでなければ、head := head の right と進めます。
• ループが終了したら false を返します。

ここで「head のキー < 0」という条件を設けているのは、リストが昇順に並んでいるため、負の値を先頭側から順に固定値として選べば、残りの正の値との組み合わせで合計0を実現できる可能性だけを効率よく調べられるからです。

C++による実装例

理解を深めるために、以下の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class TreeNode {
   public:
   int key;
   TreeNode *left;
   TreeNode *right;
   TreeNode() : key(0), left(NULL), right(NULL) {}
   TreeNode(int x) : key(x), left(NULL), right(NULL) {}
};
void bst_to_doubli_list(TreeNode* root, TreeNode** head, TreeNode** tail) {
   if (root == NULL)
      return;
   if (root->left)
      bst_to_doubli_list(root->left, head, tail);
   root->left = *tail;
   if (*tail)
      (*tail)->right = root;
   else
      *head = root;
      *tail = root;
   if (root->right)
      bst_to_doubli_list(root->right, head, tail);
}
bool is_in_double_list(TreeNode* head, TreeNode* tail, int sum) {
   while (head != tail) {
      int current = head->key + tail->key;
      if (current == sum)
         return true;
      else if (current > sum)
         tail = tail->left;
      else
         head = head->right;
   }
   return false;
}
bool is_valid_triplet(TreeNode *root) {
   if (root == NULL)
      return false;
   TreeNode* head = NULL;
   TreeNode* tail = NULL;
   bst_to_doubli_list(root, &head, &tail);
   while ((head->right != tail) && (head->key < 0)){
      if (is_in_double_list(head->right, tail, -1*head->key))
         return true;
      else
         head = head->right;
   }
   return false;
}
TreeNode* insert(TreeNode* root, int key) {
   if (root == NULL)
      return new TreeNode(key);
   if (root->key > key)
      root->left = insert(root->left, key);
   else
      root->right = insert(root->right, key);
   return root;
}
int main(){
   TreeNode* root = NULL;
   root = insert(root, 7);
   root = insert(root, -15);
   root = insert(root, 15);
   root = insert(root, -7);
   root = insert(root, 14);
   root = insert(root, 16);
   root = insert(root, 8);
   cout << is_valid_triplet(root);
}

入力

root = insert(root, 7);
root = insert(root, -15);
root = insert(root, 15);
root = insert(root, -7);
root = insert(root, 14);
root = insert(root, 16);
root = insert(root, 8);

出力

1

出力が 1（true） となっており、[-15, 7, 8] という合計が0になるトリプレットが正しく検出できていることがわかります。