Cấu trúc dữ liệu và Giải thuật Bài 77: Cây KD (K-Dimensional Tree) - Tìm kiếm láng giềng gần nhất trong không gian đa chiều.

Trong các bài học trước, chúng ta đã xử lý dữ liệu 1 chiều (mảng tuyến tính, chuỗi văn bản) vô cùng thành công. Tuy nhiên, khi xây dựng hệ thống backend cho các siêu ứng dụng như Grab, Google Maps, hoặc hệ thống Vector Database cho Trí tuệ nhân tạo (AI/LLM), dữ liệu không còn là một chuỗi ký tự mà biến thành các Tọa độ không gian (Coordinates) hoặc Vector đa chiều (N-dimensional Vectors).

Giả sử hệ thống của Hasaki cần tìm "Cửa hàng gần vị trí của người dùng nhất" dựa trên tọa độ (Vĩ độ, Kinh độ). Nếu có 10,000 cửa hàng, thuật toán vét cạn tính khoảng cách từ người dùng đến từng cửa hàng sẽ mất $O(N)$ thời gian. Khi scale lên hàng triệu shipper hoặc hàng tỷ vector nhúng (embeddings), hệ thống sẽ sụp đổ.

Để giải quyết bài toán Tìm kiếm láng giềng gần nhất (Nearest Neighbor Search) trong thời gian $O(\log N)$, chúng ta sử dụng một cấu trúc dữ liệu phân chia không gian không gian hình học: Cây KD (K-Dimensional Tree).

1. Giải phẫu Cây KD (K-Dimensional Tree)

Cây KD bản chất là một Cây tìm kiếm nhị phân (BST), nhưng thay vì so sánh một giá trị đơn lẻ, nó luân phiên thay đổi Trục tọa độ (Axis) dùng để so sánh tại mỗi độ sâu của cây.

Hãy tưởng tượng một không gian 2D (trục X và Y):

1. Tầng Root (Độ sâu 0): Chúng ta chọn trục X. Sắp xếp tất cả các điểm theo tọa độ X, chọn điểm ở giữa (Median) làm Root. Điểm này vẽ một đường thẳng dọc chia đôi không gian. Các điểm bên trái đi vào nhánh con trái, bên phải đi vào nhánh con phải.

2. Tầng 1 (Độ sâu 1): Tại các node con, chúng ta đổi sang trục Y. Chọn điểm Median theo tọa độ Y. Điểm này vẽ một đường thẳng ngang chia đôi các nửa không gian trước đó.

3. Tầng 2 (Độ sâu 2): Quay lại cắt theo trục X.

4. Quá trình này lặp lại đệ quy (Luân phiên $X \rightarrow Y \rightarrow X$ cho không gian 2D, hoặc $X \rightarrow Y \rightarrow Z \rightarrow X$ cho không gian 3D).

Mỗi Node trên cây KD không chỉ là một điểm dữ liệu, mà còn đại diện cho một Siêu mặt phẳng (Hyperplane) cắt ngang và chia không gian thành hai nửa độc lập.

2. Bí mật của Thuật toán Tìm kiếm Láng giềng (Nearest Neighbor Search)

Tìm kiếm trên KD Tree giống như việc khoanh vùng một hình tròn (hoặc khối cầu) xung quanh điểm mục tiêu (Target Point) và từ từ thu hẹp bán kính lại.

Quá trình quét diễn ra theo 2 bước:

1. Đi xuống (Traverse Down): Giống hệt BST thông thường, đi theo nhánh chứa tọa độ gần hơn với Target cho đến khi chạm đến Nút lá (Leaf). Cập nhật khoảng cách ngắn nhất tạm thời (Best Distance).

2. Quay lui (Backtracking - Bước ma thuật): Khi quay ngược lên gốc, thuật toán thực hiện một phép kiểm tra sinh tử:

   • Tính khoảng cách vuông góc từ Target đến Siêu mặt phẳng của Node hiện tại (VD: Nếu đang cắt theo trục X, khoảng cách này đơn giản là $\vert{}Target.X - Node.X\vert{}$).

   • Nếu khoảng cách này lớn hơn Best Distance hiện tại, ta tự tin bỏ qua toàn bộ nhánh con phía bên kia mặt phẳng (Pruning/Cắt tỉa) vì chắc chắn không có điểm nào gần hơn ở đó.

   • Nếu khoảng cách này nhỏ hơn, nghĩa là khối cầu tìm kiếm của ta đã "cắt qua" mặt phẳng, ta bắt buộc phải rẽ sang nhánh bên kia để tìm kiếm.

Nhờ cơ chế "Cắt tỉa" (Pruning) không gian này, thuật toán loại bỏ được 90% số điểm cần tính toán.

3. Triển khai KD-Tree bằng Golang

Dưới đây là cài đặt chuẩn mực cho một KD-Tree hỗ trợ không gian K-chiều tổng quát, kết hợp hàm xây dựng và hàm tìm kiếm điểm gần nhất:

Go

package main

import (
	"fmt"
	"math"
	"sort"
)

// Điểm trong không gian K-chiều
type Point struct {
	Coords []float64
}

// Cấu trúc Node của KD-Tree
type KDNode struct {
	Point Point
	Left  *KDNode
	Right *KDNode
	Axis  int
}

// Hàm tính khoảng cách Euclid bình phương (Tránh dùng math.Sqrt để tối ưu tốc độ CPU)
func sqDistance(p1, p2 Point) float64 {
	sum := 0.0
	for i := range p1.Coords {
		diff := p1.Coords[i] - p2.Coords[i]
		sum += diff * diff
	}
	return sum
}

// Xây dựng KD-Tree đệ quy
func BuildKDTree(points []Point, depth int) *KDNode {
	if len(points) == 0 {
		return nil
	}

	k := len(points[0].Coords)
	axis := depth % k

	// Sắp xếp các điểm theo trục hiện tại
	sort.Slice(points, func(i, j int) bool {
		return points[i].Coords[axis] < points[j].Coords[axis]
	})

	// Chọn điểm Median làm gốc
	medianIdx := len(points) / 2
	
	return &KDNode{
		Point: points[medianIdx],
		Left:  BuildKDTree(points[:medianIdx], depth+1),
		Right: BuildKDTree(points[medianIdx+1:], depth+1),
		Axis:  axis,
	}
}

// Tìm kiếm láng giềng gần nhất
func NearestNeighbor(node *KDNode, target Point, bestNode **KDNode, bestDist *float64) {
	if node == nil {
		return
	}

	// Cập nhật Best Node nếu tìm thấy điểm gần hơn
	dist := sqDistance(node.Point, target)
	if dist < *bestDist {
		*bestDist = dist
		*bestNode = node
	}

	// Quyết định nhánh ưu tiên để đi xuống sâu hơn
	var goodSide, badSide *KDNode
	axis := node.Axis
	if target.Coords[axis] < node.Point.Coords[axis] {
		goodSide = node.Left
		badSide = node.Right
	} else {
		goodSide = node.Right
		badSide = node.Left
	}

	// 1. Đi sâu xuống nhánh ưu tiên
	NearestNeighbor(goodSide, target, bestNode, bestDist)

	// 2. Backtrack: Kiểm tra xem có cần rẽ sang nhánh Bad Side hay không
	// Khoảng cách vuông góc từ Target tới mặt phẳng cắt
	planeDist := target.Coords[axis] - node.Point.Coords[axis]
	if planeDist*planeDist < *bestDist {
		NearestNeighbor(badSide, target, bestNode, bestDist)
	}
}

func main() {
	// Danh sách các tọa độ 2D
	points := []Point{
		{Coords: []float64{3, 6}},
		{Coords: []float64{17, 15}},
		{Coords: []float64{13, 15}},
		{Coords: []float64{6, 12}},
		{Coords: []float64{9, 1}},
		{Coords: []float64{2, 7}},
		{Coords: []float64{10, 19}},
	}

	// Xây dựng cây
	root := BuildKDTree(points, 0)

	// Điểm cần tìm
	target := Point{Coords: []float64{12, 14}}
	
	var bestNode *KDNode
	bestDist := math.MaxFloat64

	NearestNeighbor(root, target, &bestNode, &bestDist)

	fmt.Printf("Điểm mục tiêu: %v\n", target.Coords)
	fmt.Printf("Điểm gần nhất trên cây: %v (Khoảng cách bình phương: %.2f)\n", bestNode.Point.Coords, bestDist)
	// Mong đợi: Điểm [13, 15] là gần với [12, 14] nhất.
}

4. Phân tích Hiệu năng (Big-O) và Lời nguyền Chiều (Curse of Dimensionality)

• Thời gian Xây dựng: $O(N \log^2 N)$ nếu dùng thuật toán sort mảng cơ bản tại mỗi tầng, hoặc có thể tối ưu xuống $O(N \log N)$ bằng thuật toán Quickselect tìm Median.

• Thời gian Tìm kiếm (Search Time):

  • Trung bình: $O(\log N)$. Thuật toán cắt tỉa không gian cực kỳ hiệu quả.

  • Trường hợp xấu nhất: $O(N)$.

• Không gian (Space): $O(N)$ để lưu trữ cấu trúc Cây.

Lưu ý kỹ thuật: Cây KD tỏa sáng rực rỡ trong không gian có số chiều nhỏ (như 2D, 3D). Tuy nhiên, khi số chiều $K$ lớn (ví dụ Vector 256 chiều), hiện tượng "Lời nguyền số chiều" (Curse of Dimensionality) xuất hiện: Hầu hết mọi không gian con đều giao nhau với khối cầu bán kính tìm kiếm, khiến thuật toán phải duyệt qua gần như toàn bộ các node, làm suy thoái tốc độ về mức $O(N)$ ngang với vét cạn.

5. Ứng dụng Thực tế trong Hệ thống Hiện đại

• Geospatial Queries (Truy vấn Không gian/Địa lý): Tìm "Các tài xế xe ôm nằm trong bán kính 2km" của hệ thống Ride-hailing, hoặc quét các cụm điểm POI (Point of Interest) trên bản đồ.

• Computer Graphics & Vật lý (Ray Tracing): Trong lập trình đồ họa và game, KD-Tree (cùng biến thể BVH) phân rã không gian 3D để xác định nhanh các mặt phẳng/vật thể mà một tia sáng (Ray) sẽ chiếu trúng.

• Machine Learning (Tiền thân của Vector Database): Tìm kiếm K-Nearest Neighbors (K-NN) cho các thuật toán phân cụm dữ liệu. Trong các hệ thống RAG và LLM hiện tại, vì số chiều quá lớn, người ta thường dùng các biến thể tối ưu hơn từ gốc của KD-Tree như HNSW (Hierarchical Navigable Small World) hoặc LSH.

Xử lý không gian lưu trữ đã xong, nhưng nếu bạn cần tạo ra một hệ thống Cơ sở dữ liệu (Database) khổng lồ như MySQL hay PostgreSQL, cây nhị phân sẽ thất bại thảm hại do quá trình đọc ghi Ổ cứng (Disk I/O). Để chinh phục hạ tầng lưu trữ gốc, chúng ta cần tiến vào trái tim của mọi Database Relational.