Developer Biography

action speaks louder than word

Posted on Edited on

Problem Description

link

Problem Analysis

We can divide all people into different sets based on the languages they understand and speak. Everyone in a set can communicate with others in the same set—that is, they can understand the languages spoken by each other. Finally, we subtract the number of people in the largest set from the total number of people, which gives us the minimum number of people that need to be excluded.

Step 1: Convert Strings to IDs

All the information provided in the input appears in string form, which is extremely inconvenient for graph construction and traversal. Therefore, we first record strings and assign corresponding IDs using a dictionary and an array.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
fn build_language_dict(people: &Vec<Person>) -> HashMap<&String, usize> {
    let mut lang_dict = HashMap::new();
    let mut num_of_lang = 0usize;
    for person in people {
        for lang in person.understand.borrow() {
            if lang_dict.contains_key(lang) {
                continue;
            }
            lang_dict.insert(lang, num_of_lang);
            num_of_lang += 1;
        }
    }
    lang_dict
}
fn record_people(n: usize) -> Vec<Person> {
    let mut people: Vec<Person> = vec![];
    for _ in 0..n {
        let input = read_line_of_string();
        people.push(Person {
            speak: input[1].to_owned(),
            understand: HashSet::new(),
        });
        let id = people.len() - 1;
        for i in 1..input.len() {
            people[id].understand.insert(input[i].clone());
        }
    }
    people
}

The record_people function records each line of input, then converts them into Person structs and stores them in the people array, allowing direct access to each person via their ID later. The build_language_dict function creates a HashMap with String keys and usize values, enabling us to find the corresponding value by language name. After converting all strings to IDs, we can begin constructing the graph.

Step 2: Construct the Language Conversion Graph

For each person, they can convert the languages they understand to the language they speak. This is equivalent to drawing a directed edge from the understood language to the spoken language. By traversing each person in this way, we can build a graph of language conversion.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
fn build_lang_graph(people: &Vec<Person>, lang_dict: &HashMap<&String, usize>) -> Vec<Vec<usize>> {
    let mut lang_graph: Vec<Vec<usize>> = vec![vec![]; lang_dict.len()];
    for person in people {
        for lang in person.understand.borrow() {
            if *lang == person.speak {
                continue;
            }
            lang_graph[lang_dict[lang]].push(lang_dict[&person.speak]);
        }
    }
    for i in 0..lang_graph.len() {
        let tmp: HashSet<_> = lang_graph[i].drain(..).collect();
        lang_graph[i].extend(tmp.into_iter());
    }
    lang_graph
}

The function’s end uses a HashSet for edge deduplication. Values in the array are placed into a HashSet and then back into the array to achieve deduplication.

Step 3: Verify Two People Can Communicate

After the language graph is established, a function is needed to verify whether two people can understand each other’s spoken languages. If one person’s spoken language can be converted into one of the languages the other person understands, then it can be said that the second person can understand the first person’s speech, and vice versa. This is implemented using a naive DFS algorithm.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
fn test_connected_dfs(
    cur: usize,
    target: &HashSet<usize>,
    lang_graph: &Vec<Vec<usize>>,
    visited: &mut Vec<bool>,
) -> bool {
    if target.contains(&cur) {
        return true;
    }
    visited[cur] = true;
    for next in &lang_graph[cur] {
        if visited[*next] {
            continue;
        }
        if test_connected_dfs(*next, target, lang_graph, visited) {
            return true;
        }
    }
    false
}
fn test_biconnected(
    a: &Person,
    b: &Person,
    lang_graph: &Vec<Vec<usize>>,
    lang_dict: &HashMap<&String, usize>,
) -> bool {
    let mut visited = vec![false; lang_graph.len()];
    let target: HashSet<usize> = b
        .understand
        .borrow()
        .into_iter()
        .map(|s| lang_dict[s])
        .collect();
    if !test_connected_dfs(lang_dict[&a.speak], &target, lang_graph, &mut visited) {
       

 return false;
    }
    visited = vec![false; lang_graph.len()];
    let target: HashSet<usize> = a
        .understand
        .borrow()
        .into_iter()
        .map(|s| lang_dict[s])
        .collect();
    if !test_connected_dfs(lang_dict[&b.speak], &target, lang_graph, &mut visited) {
        return false;
    }
    true
}

The test_biconnected function encapsulates test_connected_dfs. By calling this subfunction twice, it verifies whether two people can communicate.

Step 4: Construct a Disjoint Set Union

When analyzing the relationships between each person, a disjoint set union (DSU) data structure is used for organization. Everyone in a DSU can communicate with each other, and if two people from different DSUs can communicate, those DSUs are merged.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
struct DisjointSetUnion {
    fa: Vec<usize>,
}
impl DisjointSetUnion {
    fn get_fa(&mut self, cur: usize) -> usize {
        if self.fa[cur] == cur {
            return cur;
        }
        let tmp = self.fa[cur];
        self.fa[cur] = self.get_fa(tmp);
        self.fa[cur]
    }
}
fn build_people_disjoint_set_union(
    people: &Vec<Person>,
    lang_graph: &Vec<Vec<usize>>,
    lang_dict: &HashMap<&String, usize>,
) -> DisjointSetUnion {
    let mut people_disjoint_set_union = DisjointSetUnion {
        fa: (0..people.len()).into_iter().collect(),
    };
    for i in 0..people.len() {
        for j in i + 1..people.len() {
            if people_disjoint_set_union.get_fa(i) == people_disjoint_set_union.get_fa(j)
                || !test_biconnected(&people[i], &people[j], lang_graph, lang_dict)
            {
                continue;
            }
            let tmp = people_disjoint_set_union.get_fa(i);
            people_disjoint_set_union.fa[tmp] = people_disjoint_set_union.get_fa(j);
        }
    }
    people_disjoint_set_union
}

Path compression is used in the DSU implementation. Each call to get_fa updates the node’s parent to the root node of the set, thereby flattening the DSU and optimizing query efficiency.

Step 5: Obtain the Largest DSU Count

An array records the count of each DSU, and each node’s DSU is identified using get_fa, updating the count accordingly. The function returns the count of the largest DSU.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
impl DisjointSetUnion {
    fn get_fa(&mut self, cur: usize) -> usize {
        if self.fa[cur] == cur {
            return cur;
        }
        let tmp = self.fa[cur];
        self.fa[cur] = self.get_fa(tmp);
        self.fa[cur]
    }
    fn get_largest_set(&mut self) -> usize {
        let mut num_of_set = vec![0usize; self.fa.len()];
        for i in 0..self.fa.len() {
            num_of_set[self.get_fa(i)] += 1;
        }
        *num_of_set.iter().max().unwrap()
    }
}

Step 6: Integrate and Return the Answer

Arrange all subfunctions in the main function and print the value of n - people_disjoint_set_union.get_largest_set() as the result.

1
2
3
4
5
6
7
8
9
fn main() {
    let n = read_line_of_i32()[0] as usize;
    let people = record_people(n);
    let lang_dict = build_language_dict(&people);
    let lang_graph = build_lang_graph(&people, &lang_dict);
    let mut people_disjoint_set_union =
        build_people_disjoint_set_union(&people, &lang_graph, &lang_dict);
    println!("{}", n - people_disjoint_set_union.get_largest_set());
}

Conclusion: Complete Code

The provided code demonstrates a method for solving the problem by dividing people into sets based on language understanding and speaking, using data structures and algorithms such as hash maps, graphs, DFS, and disjoint set unions to efficiently find the minimum number of people that need to be excluded to allow for universal communication within the remaining group.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
use std::{
    borrow::Borrow,
    collections::{HashMap, HashSet},
};
struct Person {
    speak: String,
    understand: HashSet<String>,
}
struct DisjointSetUnion {
    fa: Vec<usize>,
}
impl DisjointSetUnion {
    fn get_fa(&mut self, cur: usize) -> usize {
        if self.fa[cur] == cur {
            return cur;
        }
        let tmp = self.fa[cur];
        self.fa[cur] = self.get_fa(tmp);
        self.fa[cur]
    }
    fn get_largest_set(&mut self) -> usize {
        let mut num_of_set = vec![0usize; self.fa.len()];
        for i in 0..self.fa.len() {
            num_of_set[self.get_fa(i)] += 1;
        }
        *num_of_set.iter().max().unwrap()
    }
}
fn main() {
    let n = read_line_of_i32()[0] as usize;
    let people = record_people(n);
    let lang_dict = build_language_dict(&people);
    let lang_graph = build_lang_graph(&people, &lang_dict);
    let mut people_disjoint_set_union =
        build_people_disjoint_set_union(&people, &lang_graph, &lang_dict);
    println!("{}", n - people_disjoint_set_union.get_largest_set());
}
fn build_people_disjoint_set_union(
    people: &Vec<Person>,
    lang_graph: &Vec<Vec<usize>>,
    lang_dict: &HashMap<&String, usize>,
) -> DisjointSetUnion {
    let mut people_disjoint_set_union = DisjointSetUnion {
        fa: (0..people.len()).into_iter().collect(),
    };
    for i in 0..people.len() {
        for j in i + 1..people.len() {
            if people_disjoint_set_union.get_fa(i) == people_disjoint_set_union.get_fa(j)
                || !test_biconnected(&people[i], &people[j], lang_graph, lang_dict)
            {
                continue;
            }
            let tmp = people_disjoint_set_union.get_fa(i);
            people_disjoint_set_union.fa[tmp] = people_disjoint_set_union.get_fa(j);
        }
    }
    people_disjoint_set_union
}
fn test_connected_dfs(
    cur: usize,
    target: &HashSet<usize>,
    lang_graph: &Vec<Vec<usize>>,
    visited: &mut Vec<bool>,
) -> bool {
    if target.contains(&cur) {
        return true;
    }
    visited[cur] = true;
    for next in &lang_graph[cur] {
        if visited[*next] {
            continue;
        }
        if test_connected_dfs(*next, target, lang_graph, visited) {
            return true;
        }
    }
    false
}
fn test_biconnected(
    a: &Person,
    b: &Person,
    lang_graph: &Vec<Vec<usize>>,
    lang_dict: &HashMap<&String, usize>,
) -> bool {
    let mut visited = vec![false; lang_graph.len()];
    let target: HashSet<usize> = b
        .understand
        .borrow()
        .into_iter()
        .map(|s| lang_dict[s])
        .collect();
    if !test_connected_dfs(lang_dict[&a.speak], &target, lang_graph, &mut visited) {
        return false;
    }
    visited = vec![false; lang_graph.len()];
    let target: HashSet<usize> = a
        .understand
        .borrow()
        .into_iter()
        .map(|s| lang_dict[s])
        .collect();
    if !test_connected_dfs(lang_dict[&b.speak], &target, lang_graph, &mut visited) {
        return false;
    }
    true
}
fn build_lang_graph(people: &Vec<Person>, lang_dict: &HashMap<&String, usize>) -> Vec<Vec<usize>> {
    let mut lang_graph: Vec<Vec<usize>> = vec![vec![]; lang_dict.len()];
    for person in people {
        for lang in person.understand.borrow() {
            if *lang == person.speak {
                continue;
            }
            lang_graph[lang_dict[lang]].push(lang_dict[&person.speak]);
        }
    }
    for i in 0..lang_graph.len() {
        let tmp: HashSet<_> = lang_graph[i].drain(..).collect();
        lang_graph[i].extend(tmp.into_iter());
    }
    lang_graph
}
fn build_language_dict(people: &Vec<Person>) -> HashMap<&String, usize> {
    let mut lang_dict = HashMap::new();
    let mut num_of_lang = 0usize;
    for person in people {
        for lang in person.understand.borrow() {
            if lang_dict.contains_key(lang) {
                continue;
            }
            lang_dict.insert(lang, num_of_lang);
            num_of_lang += 1;
        }
    }
    lang_dict
}
fn record_people(n: usize) -> Vec<Person> {
    let mut people: Vec<Person> = vec![];
    for _ in 0..n {
        let input = read_line_of_string();
        // println!("{:?}", input);
        people.push(Person {
            speak: input[1].to_owned(),
            understand: HashSet::new(),
        });
        let id = people.len() - 1;
        for i in 1..input.len() {
            people[id].understand.insert(input[i].clone());
        }
    }
    people
}
fn read_line_of_i32() -> Vec<i32> {
    let mut buffer = String::new();
    let _ = std::io::stdin().read_line(&mut buffer);
    buffer
        .trim()
        .to_owned()
        .split(" ")
        .map(|x| x.to_owned().parse::<i32>().unwrap())
        .collect()
}
fn read_line_of_string() -> Vec<String> {
    let mut buffer = String::new();
    let _ = std::io::stdin().read_line(&mut buffer);
    buffer
        .trim()
        .to_owned()
        .split(" ")
        .map(|x| x.to_owned())
        .collect()
}

Posted on Edited on

Dynamic Programming

In both contexts it refers to simplifying a complicated problem by breaking it down into simpler sub-problems in a recursive manner. While some decision problems cannot be taken apart this way, decisions that span several points in time do often break apart recursively. Likewise, in computer science, if a problem can be solved optimally by breaking it into sub-problems and then recursively finding the optimal solutions to the sub-problems, then it is said to have optimal substructure.
If sub-problems can be nested recursively inside larger problems, so that dynamic programming methods are applicable, then there is a relation between the value of the larger problem and the values of the sub-problems. In the optimization literature this relationship is called the Bellman equation. –Wikipedia

I’m not sure how much you understood from Wikipedia’s description, but I think starting with a simple example would be better.

The Fibonacci Sequence: A Gateway to Dynamic Programming

At its core, the Fibonacci sequence is a series of numbers where each number is the sum of the two preceding ones, starting from 0 and 1. Mathematically, it’s defined as:

$$
F_i = F_{i-1} + F_{i-2} (i >= 2)\
F_1 = 1, F_0 = 0
$$

This simple formula lays the foundation for understanding the three pillars of dynamic programming: State, State Transition Equation, Boundary

State

In dynamic programming, the ‘state’ encapsulates the current situation in our problem-solving context. For the Fibonacci sequence, the state is represented by $F_i$, indicating the i-th number in the sequence. Understanding the state is crucial because it defines what we’re trying to compute.

State Transition Equation

The state transition equation describes how to transition from one state to the next. For the Fibonacci sequence, this is given by $F_i = F_{i-1} + F_{i-2}$. It shows how the current state (F_i) is derived from the previous two states ($F_{i-1}$ and $F_{i-2}$). This equation is the heart of dynamic programming, guiding how we break down problems into smaller, manageable pieces.

Boundary

Boundary conditions provide the starting points or base cases for our computation. For the Fibonacci sequence, these are $F_1$ = 1 and $F_0$ = 0. They’re essential because they prevent infinite recursion and ensure that our dynamic programming algorithm has a clear and definitive starting point.

Concepts end here. Let’s start to code.

Implementation of Fibonacci Sequence

1
2
3
4
5
6
7
def main():
    n = int(input())
    fib = [0] * (n + 2)
    fib[1] = 1
    for i in range(2, n + 2):
        fib[i] = fib[i - 1] + fib[i - 2]
    print(fib[n])

Here, we’ve crafted a function named main that computes the nth Fibonacci number. This function prompts for an input, n, and initializes an array fib with a length of n + 2 (extra two is just for safety) to store the Fibonacci sequence. It sets the second element (index 1) to 1, reflecting the sequence’s start. Through a loop, it then calculates each Fibonacci number by summing the two preceding numbers and stores them in the array. Finally, it prints the nth Fibonacci number.

Let’s execute this function:

1
2
3
35
9227465
Execution time: 1.889999839477241e-05 seconds

We’ve successfully solved our first dynamic programming problem with this implementation. However, the code can be streamlined for better readability. Let’s work on making it more intuitive.

Implement Fibonacci Recursively

1
2
3
4
5
6
def fib(i):
    return i if i < 2 else fib(i - 1) + fib(i - 2)

def main():
    n = int(input())
    print(fib(n))

This recursive implementation is more intuitive and direct, using just a single line of code to define the Fibonacci computation. Let’s execute this function:

1
2
3
35
9227465
The fib function took 2.1213817596435547 seconds to execute.

Okay, everything looks fine… Really?
Let’s see the third line of output.

1
The fib function took 2.1213817596435547 seconds to execute.

The function runs extremely slow. What happened there? Let’s check how many times does function fib() be called.

1
2
3
4
35
9227465
The fib function took 3.8617589473724365 seconds to execute.
The fib function was called 29860703 times.

The observed slowdown in executing the recursive Fibonacci function, taking over 2 seconds to compute fib(35) and resulting in an astonishing 29,860,703 function calls, highlights a critical inefficiency inherent in the naive recursive approach. This inefficiency is primarily due to the excessive number of redundant calls for calculating the same Fibonacci numbers multiple times.

To understand why so many calls are made, consider how the recursive function works for fib(n): it calls fib(n-1) and fib(n-2). Both of these calls then recursively invoke further fib calculations, leading to a combinatorial explosion of function calls. Many of these calls are repeated; for instance, fib(2) is calculated multiple times across different branches of the recursion tree.

Now let’s try to solve the problem. To address this inefficiency, a technique known as memoization can be applied. Memoization is a form of caching where the results of expensive function calls are stored, so that if the same inputs occur again, the function can return the stored result instead of recalculating it. Applying memoization to the Fibonacci sequence means that once we calculate fib(n), we store its value. If fib(n) is needed again, we simply retrieve the stored value instead of recalculating it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
cache = []

def fib(i):
    global cache
    if cache[i] != 0:
        return cache[i]
    cache[i] = i if i < 2 else fib(i - 1) + fib(i - 2)
    return cache[i]

def main():
    global cache
    n = int(input())
    cache = [0] * (n + 1)
    result = fib(n)
    print(result)

Let’s run and check the performance.

1
2
3
35
9227465
The fib function took 1.539999993838137e-05 seconds to execute.

This performance boost is due to the drastic reduction in the number of calls to the fib function, as each Fibonacci number is calculated only once. Now it has the same runtime as the method use array. But can we make it better?

Rolling array

Let’s revisit our initial dynamic programming approach to computing the nth Fibonacci number. In this version, we allocate an array large enough to hold all Fibonacci numbers up to n. However, we notice a key optimization point: to calculate any Fibonacci number, we only ever need the last two numbers in the sequence.

1
2
3
4
5
6
7
def main():
    n = int(input())
    fib = [0] * (n + 2)
    fib[1] = 1
    for i in range(2, n + 2):
        fib[i] = fib[i - 1] + fib[i - 2]
    print(fib[n])

We can find that when we try to get n-th Fibonacci Number. It will create an array with n variables. But each time we only use fib[i],fib[i - 1], and fib[i - 2]. So we can track only these three variables.

1
2
3
4
5
6
7
8
9
def main():
    n = int(input())
    fib_i = 1
    fib_i_prev = 0
    start_time = time.perf_counter()
    for i in range(2, n + 1):
        fib_i, fib_i_prev = fib_i + fib_i_prev, fib_i
    end_time = time.perf_counter()
    print(fib_i)

By maintaining just two variables, fib_i and fib_i_prev, and updating them in each iteration, we significantly reduce the space complexity of our program to O(1), while still retaining a time complexity of O(n). This rolling array technique optimizes our solution, making it both space-efficient and fast for calculating Fibonacci numbers. But what happened if it is still too slow.

Matrix Acceleration

At first glance, achieving $O(1)$ space complexity and $O(n)$ time complexity might seem like the pinnacle of optimization for any algorithm. This might lead one to prematurely conclude that there’s no further scope for enhancement. However, this article aims to challenge that notion by introducing a mathematical strategy capable of significantly boosting our program’s efficiency, especially for large-scale problems.

1
2
3
100000000
908460138
Execution time: 10.320621700004267 seconds

Consider when n scales to $10 ^ 8$. At this magnitude, an algorithm with $O(n)$ time complexity begins to show its limitations, as evidenced by an execution time exceeding 10 seconds. Such scenarios demand a more sophisticated approach to reduce the computational complexity further.

Understanding the Fibonacci Sequence with Matrix Exponentiation

The Fibonacci sequence is a classic example where exponential growth can cause simple recursive or iterative solutions to falter due to their linear or pseudo-linear time complexities. To compute the n-th Fibonacci number efficiently for very large n, we employ matrix exponentiation, a method that allows us to achieve a time complexity of $O(log_2{n})$.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
MOD = 1000000007

def main():
    n = int(input())
    print(fib(n))

def fib(n):
    if n < 2:
        return n
    st = [[1, 0, 0]]
    update_matrix = qpow([[1, 1, 0], [1, 0, 1], [0, 0, 0]], n - 1)
    return matrix_multiply(st, update_matrix)[0][0] % MOD 

def qpow(a, n):
    if n == 1:
        return a
    if n % 2 == 0:
        return qpow(matrix_multiply(a, a), n // 2)
    return matrix_multiply(a, qpow(matrix_multiply(a, a), n // 2))

def matrix_multiply(a, b):
    ans = [[0] * len(b[0]) for _ in range(len(a))]
    for i in range(len(a)):
        for j in range(len(b[0])):
            for k in range(len(b)): 
                ans[i][j] += a[i][k] * b[k][j] % MOD
    return ans

if __name__ == "__main__":
    main()

Implementation Details

The provided Python code showcases an efficient way to compute large Fibonacci numbers modulo MOD = 1000000007, ensuring that the result stays within integer limits suitable for many computational environments.

  • Matrix Power Quickening (qpow): The heart of this optimization lies in the qpow function, which implements fast exponentiation for matrices. Instead of naively multiplying the matrix n times, this function recursively squares the matrix, reducing the exponent by half each time, effectively turning the time complexity from $O(n)$ to $O(log_2{n})$.

  • Matrix Multiplication (matrix_multiply): This function multiplies two matrices, a crucial operation in our exponentiation process. Matrix multiplication is used both in the exponentiation process and in calculating the final Fibonacci number from the base case matrix.

  • Computing Fibonacci (fib): The fib function initializes a state matrix representing the base case of Fibonacci calculation. It then calculates the power of the transition matrix, which encodes the Fibonacci recurrence relation, using the optimized qpow method. The result is a matrix that, when multiplied by the state matrix, yields the n-th Fibonacci number.

  • Generalization: This method is not limited to Fibonacci numbers; it can be generalized to other problems that can be expressed in terms of matrix exponentiation, showcasing the versatility of the technique.

Posted on Edited on In

Abstraction

This Python project is designed to automate chat responses in the popular messaging application, WeChat (微信), using OpenAI’s GPT-3 model. The project leverages the power of GPT-3’s language understanding capabilities to generate contextually relevant responses to incoming messages.

Project Design

The project is structured around two main functions: the GPT function, which interacts with the GPT-3 model, and the main function, which handles the user interface and message processing in WeChat.

GPT Function

The GPT function is designed to generate responses using the GPT-3 model. It takes as input a string text representing the incoming message and a list record that maintains the conversation history. The function appends the incoming message to the conversation record, sets the OpenAI API key, and creates a completion using the OpenAI API. The completion is a response generated by the GPT-3 model based on the conversation history. The function then extracts the assistant’s message from the response, appends it to the conversation record, and returns it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# Define a function to handle the interaction with the GPT-3 model
def GPT(text: str, record: List[dict]) -> str:
    # Append the user's message to the conversation record
    record.append({"role": "user", "content": text})
    # Print the user's message for debugging purposes
    print(text)
    # Set the OpenAI API key
    openai.api_key = apiKey
    # Define the model ID
    id = "gpt-3.5-turbo"
    # Initialize a question string with the user's text
    question = "###" + text + "###"
    # Create a completion using the OpenAI API
    response = openai.ChatCompletion.create(
        model=id,
        messages=record,
        temperature=0.1,
        top_p=0.1,
        max_tokens=2048,
    )
    # Extract the total token usage from the response for debugging purposes
    usage = response.usage["total_tokens"]
    # Extract the content from the assistant's message in the response
    response = response["choices"][0]["message"]["content"]
    # Append the assistant's message to the conversation record
    record.append({"role": "assistant", "content": response})
    # Return the assistant's response
    return response

Main Function

The main function is the core of the project. It initializes the conversation with a user’s template and creates a WindowControl object for the WeChat window. It then creates a ListControl object for the chat list in the window and loads reply data from a CSV file.

The function then enters a loop where it continuously checks for unread messages in the chat. When an unread message is found, it clicks on it, extracts the content of the last message, and generates a reply using the GPT function. If the conversation record has more than six messages, it trims it down to maintain a manageable size. The function then prepares the reply to be sent. If there is a reply, it sends it; if there is no reply, it sends a default message.

The script is designed to handle exceptions gracefully. If an error occurs during the execution of the main function, it prints the error and breaks the loop.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def main():
    # Initialize the conversation with the user's template
    mem = [{"role": "user", "content": template}]
    try:
        # Create a WindowControl object for the "微信" window
        wx = WindowControl(Name="微信")
        print(wx)
        wx.SwitchToThisWindow()
        # Create a ListControl object for the "会话" list in the window
        chat = wx.ListControl(Name="会话")
        print("Find chat and combine:", chat)
    except Exception as e:
        print(f"Error: {e}")
        return

    # Main loop
    while True:
        try:
            # Look for unread messages in the chat
            unread = chat.TextControl(searchDepth=4)
            # Wait until an unread message appears
            while not unread.Exists(0):
                pass
            print("Find unread:", unread)
            # If there is an unread message, click on it
            if not unread.Name: continue
            unread.Click(simulateMove=False)
            # Extract the content of the last message
            last_msg = wx.ListControl(Name="消息").GetChildren()[-1].Name
            print("Last message:", last_msg)
            # Generate a reply using the GPT-3 model
            reply = GPT(last_msg, mem)
            # If the conversation record has more than 6 messages, trim it down
            if len(mem) > 6:
                mem = mem[3:]
                mem = [{"role": "user", "content": template}] + mem
            print("memory size:", len(mem))
            # Prepare the reply to be sent
            ar = [reply]
            # If there is a reply, send it
            if ar:
                wx.SwitchToThisWindow()
                wx.SendKeys(ar[0].replace(
                    "{br}", "{Shift}{Enter}"), waitTime=0)
                wx.SendKeys("{Enter}", waitTime=0)
                wx.TextControl(SubName=ar[0][:5]).Click()
            else:
                # If there is no reply, send a default message
                wx.SwitchToThisWindow()
                wx.SendKeys("Unknown", waitTime=0)
                wx.SendKeys("{Enter}", waitTime=0)
                wx.TextControl(SubName=last_msg[:5]).Click()
        except Exception as e:
            print(f"Error: {e}")
            break

Execution

The script is designed to be run directly. If the script is run directly (i.e., not imported as a module), it calls the main function and starts the automated chat response process.

Conclusion

In summary, this Python project is a sophisticated example of how to use OpenAI’s GPT-3 model to automate responses in a chat application. It demonstrates how to interact with the GPT-3 model, handle user interface elements in a chat application, and manage conversation history for context-aware responses. The project is a testament to the potential of AI in automating tasks and enhancing user experience in messaging platforms.

Complete Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
from typing import List
import pandas as pd
import numpy as np
from uiautomation import WindowControl
import openai

apiKey = ""

template = ""

# Define a function to handle the interaction with the GPT-3 model
def GPT(text: str, record: List[dict]) -> str:
    # Append the user's message to the conversation record
    record.append({"role": "user", "content": text})
    # Print the user's message for debugging purposes
    print(text)
    # Set the OpenAI API key
    openai.api_key = apiKey
    # Define the model ID
    id = "gpt-3.5-turbo"
    # Initialize a question string with the user's text
    question = "###" + text + "###"
    # Create a completion using the OpenAI API
    response = openai.ChatCompletion.create(
        model=id,
        messages=record,
        temperature=0.1,
        top_p=0.1,
        max_tokens=2048,
    )
    # Extract the total token usage from the response for debugging purposes
    usage = response.usage["total_tokens"]
    # Extract the content from the assistant's message in the response
    response = response["choices"][0]["message"]["content"]
    # Append the assistant's message to the conversation record
    record.append({"role": "assistant", "content": response})
    # Return the assistant's response
    return response


# Main function
def main():
    # Initialize the conversation with the user's template
    mem = [{"role": "user", "content": template}]
    try:
        # Create a WindowControl object for the "微信" window
        wx = WindowControl(Name="微信")
        print(wx)
        wx.SwitchToThisWindow()
        # Create a ListControl object for the "会话" list in the window
        chat = wx.ListControl(Name="会话")
        print("Find chat and combine:", chat)
    except Exception as e:
        print(f"Error: {e}")
        return

    # Main loop
    while True:
        try:
            # Look for unread messages in the chat
            unread = chat.TextControl(searchDepth=4)
            # Wait until an unread message appears
            while not unread.Exists(0):
                pass
            print("Find unread:", unread)
            # If there is an unread message, click on it
            if not unread.Name: continue
            unread.Click(simulateMove=False)
            # Extract the content of the last message
            last_msg = wx.ListControl(Name="消息").GetChildren()[-1].Name
            print("Last message:", last_msg)
            # Generate a reply using the GPT-3 model
            reply = GPT(last_msg, mem)
            # If the conversation record has more than 6 messages, trim it down
            if len(mem) > 6:
                mem = mem[3:]
                mem = [{"role": "user", "content": template}] + mem
            print("memory size:", len(mem))
            # Prepare the reply to be sent
            ar = [reply]
            # If there is a reply, send it
            if ar:
                wx.SwitchToThisWindow()
                wx.SendKeys(ar[0].replace(
                    "{br}", "{Shift}{Enter}"), waitTime=0)
                wx.SendKeys("{Enter}", waitTime=0)
                wx.TextControl(SubName=ar[0][:5]).Click()
            else:
                # If there is no reply, send a default message
                wx.SwitchToThisWindow()
                wx.SendKeys("Unknown", waitTime=0)
                wx.SendKeys("{Enter}", waitTime=0)
                wx.TextControl(SubName=last_msg[:5]).Click()
        except Exception as e:
            print(f"Error: {e}")
            break


# Run the main function if this script is run directly
if __name__ == "__main__":
    main()

Posted on Edited on In

English Version

摘要

本文介绍了一种功能强大的翻译工具，将 OpenAI 的先进 GPT-3.5 Turbo 语言模型与用户友好的 Wox 启动器插件系统相结合。该工具可以识别输入文本的语言，将中文文本翻译成英文或将英文文本翻译成中文，并用中文提供解释和示例。本文阐述了核心功能，与 Wox 启动器的交互以及使用 pyperclip 库实现的剪贴板集成。我们还提供了一段代码示例，以演示其实现。

引言

在本文中，我们介绍了一种功能强大的翻译工具，将 OpenAI 的尖端 GPT-3.5 Turbo 语言模型与用户友好的 Wox 启动器插件系统相结合。该工具旨在识别输入文本的语言，将中文文本翻译成英文或其他语言的文本翻译成中文，并提供中文解释和示例。

利用 GPT-3.5 Turbo

OpenAI 的 GPT-3.5 Turbo 语言模型以高效处理复杂语言任务而闻名。该翻译工具利用模型的功能，通过实现一个翻译函数来接收文本输入，构建一个格式化的问题发送给模型，并接收响应。然后解析响应，提取相关信息并将其作为列表返回。您可以在此处获取有关 ChatAPI 的更多信息：openAI 文档

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def translate(text: str) -> list:
    openai.api_key = KEY
    id = "gpt-3.5-turbo"
    question = """
1: Check the text's language
2:  If language != Chinese {
        translate it into Chinese 
    } else {
        translate it into English
    } 
3: Desired format:

language: -||-
translate: -||-
explanation: -||-
example: -||-
##

Text1: ###Hello###
language: English
translate: 你好
explanation: 这是常用的问候语
example: Hello, how are you today?
##

Text2: ###你好###
language: Chinese
translate: Hello
explanation: 这是常用的问候语
example: 你好，你今天怎么样?
##

Text3: ###こんにちは###
language: Japanese
translate: 你好
explanation: 这是常用的问候语
example: こんにちは、今日はお元気ですか?
##

Text4: ###안녕하세요###
language: Korean
translate: 你好
explanation: 这是常用的问候语
example: 안녕하세요, 오늘은 어떻게 지내세요?
##

Text5: """
    question += "###" + text + "###"
    response = openai.ChatCompletion.create(
        model=id,
        messages=[{
            "role": "user",
            "content": question
        }],
        temperature=0,
        top_p=0.1,
        max_tokens=2048,
    )
    usage = response.usage["total_tokens"]
    response = response["choices"][0]["message"]["content"]
    # print(response)
    response = list(response.split("\n"))
    response.append("usage: " + str(usage))
    return response

与 Wox 启动器集成

Wox 启动器插件系统为用户提供了一个方便的互动的界面来使用翻译工具。通过创建一个继承自 Wox 类的主类，该工具定义了一个查询方法来处理用户输入并显示结果。当用户输入一个以 “&” 结尾的关键字时，代码调用翻译函数，将响应显示为可选择的项目列表。您可以在此处获取有关 Wox 启动器插件开发的更多信息：Wox 文档

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class main(Wox):
    def query(self, keyword: str):
        results = []
        if not keyword.endswith('&'):
            return results
        response: list[str] = translate(keyword[:-1])
        for reponse in response:
            results.append({
                "IcoPath": "Images/ico.png",
                "Title": reponse,
                "SubTitle": "Copy",
                "JsonRPCAction": {
                    "method": "copy",
                    "parameters": [reponse],
                    "dontHideAfterAction": False
                },
            })
        return results

    def copy(self, word: str):
        pyperclip.copy(word)

集成剪贴板

为了提高用户体验，脚本集成了 pyperclip 库，使用户只需点击即可将所选文本复制到剪贴板。在选择结果时，触发复制方法，轻松为用户复制文本以便粘贴到其他应用程序中。

显示

Display

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import openai
from wox import Wox
import pyperclip

KEY = ""

def translate(text: str) -> list:
    openai.api_key = KEY
    id = "gpt-3.5-turbo"
    question = """
1: Check the text's language
2:  If language != Chinese {
        translate it into Chinese 
    } else {
        translate it into English
    } 
3: Desired format:

language: -||-
translate: -||-
explanation: -||-
example: -||-
##

Text1: ###Hello###
language: English
translate: 你好
explanation: 这是常用的问候语
example: Hello, how are you today?
##

Text2: ###你好###
language: Chinese
translate: Hello
explanation: 这是常用的问候语
example: 你好，你今天怎么样?
##

Text3: ###こんにちは###
language: Japanese
translate: 你好
explanation: 这是常用的问候语
example: こんにちは、今日はお元気ですか?
##

Text4: ###안녕하세요###
language: Korean
translate: 你好
explanation: 这是常用的问候语
example: 안녕하세요, 오늘은 어떻게 지내세요?
##

Text5: """
    question += "###" + text + "###"
    response = openai.ChatCompletion.create(
        model=id,
        messages=[{
            "role": "user",
            "content": question
        }],
        temperature=0,
        top_p=0.1,
        max_tokens=2048,
    )
    usage = response.usage["total_tokens"]
    response = response["choices"][0]["message"]["content"]
    # print(response)
    response = list(response.split("\n"))
    response.append("usage: " + str(usage))
    return response


class main(Wox):
    def query(self, keyword: str):
        results = []
        if not keyword.endswith('&'):
            return results
        response: list[str] = translate(keyword[:-1])
        for reponse in response:
            results.append({
                "IcoPath": "Images/ico.png",
                "Title": reponse,
                "SubTitle": "Copy",
                "JsonRPCAction": {
                    "method": "copy",
                    "parameters": [reponse],
                    "dontHideAfterAction": False      # 运行后是否隐藏Wox窗口
                },
            })
        return results

    def copy(self, word: str):
        pyperclip.copy(word)


if __name__ == "__main__":
    main()

Posted on Edited on In

中文版

Abstract

This article presents a powerful translation tool that combines the advanced GPT-3.5 Turbo language model by OpenAI with the user-friendly Wox launcher plugin system. The tool identifies the input text’s language, translates Chinese text to English or English text to Chinese, and provides an explanation and example in Chinese. The article explains the core functionality, interaction with Wox launcher, and clipboard integration using the pyperclip library. We also provide an example of the code to demonstrate its implementation.

Introduction

In this article, we introduce a powerful translation tool that integrates the cutting-edge GPT-3.5 Turbo language model by OpenAI with the user-friendly Wox launcher plugin system. This tool is designed to identify the input text’s language, translate Chinese text to English or other languages’ text to Chinese, and provide an explanation and an example in Chinese.

Harnessing GPT-3.5 Turbo

The GPT-3.5 Turbo language model by OpenAI is known for its ability to handle complex language tasks efficiently. This translation tool utilizes the model’s capabilities by implementing a translate function that takes a text input, constructs a formatted question to send to the model, and receives a response. The response is then parsed, extracting the relevant information and returning it as a list. You can get more information about ChatAPI here: openAI documentation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def translate(text: str) -> list:
    openai.api_key = KEY
    id = "gpt-3.5-turbo"
    question = """
1: Check the text's language
2:  If language != Chinese {
        translate it into Chinese 
    } else {
        translate it into English
    } 
3: Desired format:

language: -||-
translate: -||-
explanation: -||-
example: -||-
##

Text1: ###Hello###
language: English
translate: 你好
explanation: 这是常用的问候语
example: Hello, how are you today?
##

Text2: ###你好###
language: Chinese
translate: Hello
explanation: 这是常用的问候语
example: 你好，你今天怎么样?
##

Text3: ###こんにちは###
language: Japanese
translate: 你好
explanation: 这是常用的问候语
example: こんにちは、今日はお元気ですか?
##

Text4: ###안녕하세요###
language: Korean
translate: 你好
explanation: 这是常用的问候语
example: 안녕하세요, 오늘은 어떻게 지내세요?
##

Text5: """
    question += "###" + text + "###"
    response = openai.ChatCompletion.create(
        model=id,
        messages=[{
            "role": "user",
            "content": question
        }],
        temperature=0,
        top_p=0.1,
        max_tokens=2048,
    )
    usage = response.usage["total_tokens"]
    response = response["choices"][0]["message"]["content"]
    # print(response)
    response = list(response.split("\n"))
    response.append("usage: " + str(usage))
    return response

Integrating with Wox Launcher

The Wox launcher plugin system offers a convenient and interactive interface for users to engage with the translation tool. By creating a main class that inherits from the Wox class, the tool defines a query method to handle user input and display results. When a user enters a keyword ending with an ampersand (&), the code calls the translate function, displaying the response as a list of selectable items. You can get more information about the Wox launcher plugin development here: Wox Documentation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class main(Wox):
    def query(self, keyword: str):
        results = []
        if not keyword.endswith('&'):
            return results
        response: list[str] = translate(keyword[:-1])
        for reponse in response:
            results.append({
                "IcoPath": "Images/ico.png",
                "Title": reponse,
                "SubTitle": "Copy",
                "JsonRPCAction": {
                    "method": "copy",
                    "parameters": [reponse],
                    "dontHideAfterAction": False
                },
            })
        return results

    def copy(self, word: str):
        pyperclip.copy(word)

Clipboard Integration for User Convenience

To enhance user experience, the script incorporates the pyperclip library, enabling users to copy the selected text to the clipboard with just a click. Upon selecting a result, the copy method is triggered, effortlessly copying the text for users to paste into other applications.

Display

Display

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import openai
from wox import Wox
import pyperclip

KEY = ""

def translate(text: str) -> list:
    openai.api_key = KEY
    id = "gpt-3.5-turbo"
    question = """
1: Check the text's language
2:  If language != Chinese {
        translate it into Chinese 
    } else {
        translate it into English
    } 
3: Desired format:

language: -||-
translate: -||-
explanation: -||-
example: -||-
##

Text1: ###Hello###
language: English
translate: 你好
explanation: 这是常用的问候语
example: Hello, how are you today?
##

Text2: ###你好###
language: Chinese
translate: Hello
explanation: 这是常用的问候语
example: 你好，你今天怎么样?
##

Text3: ###こんにちは###
language: Japanese
translate: 你好
explanation: 这是常用的问候语
example: こんにちは、今日はお元気ですか?
##

Text4: ###안녕하세요###
language: Korean
translate: 你好
explanation: 这是常用的问候语
example: 안녕하세요, 오늘은 어떻게 지내세요?
##

Text5: """
    question += "###" + text + "###"
    response = openai.ChatCompletion.create(
        model=id,
        messages=[{
            "role": "user",
            "content": question
        }],
        temperature=0,
        top_p=0.1,
        max_tokens=2048,
    )
    usage = response.usage["total_tokens"]
    response = response["choices"][0]["message"]["content"]
    # print(response)
    response = list(response.split("\n"))
    response.append("usage: " + str(usage))
    return response


class main(Wox):
    def query(self, keyword: str):
        results = []
        if not keyword.endswith('&'):
            return results
        response: list[str] = translate(keyword[:-1])
        for reponse in response:
            results.append({
                "IcoPath": "Images/ico.png",
                "Title": reponse,
                "SubTitle": "Copy",
                "JsonRPCAction": {
                    "method": "copy",
                    "parameters": [reponse],
                    "dontHideAfterAction": False      # 运行后是否隐藏Wox窗口
                },
            })
        return results

    def copy(self, word: str):
        pyperclip.copy(word)


if __name__ == "__main__":
    main()

Posted on Edited on In

English Version

摘要

红黑树是一种自平衡的二叉搜索树，用于按排序顺序存储和检索元素，并具有O（log n）时间复杂度的所有主要操作。它们被设计为内存高效和平衡，使它们比其他平衡树如AVL树更高效。在本文中，我们讨论了在Rust中实现红黑树，这是一种为性能和安全性而设计的现代编程语言。我们讨论了树的结构，以及插入，删除和搜索的函数。此外，我们还讨论了修复函数，用于在每次插入和删除后恢复树的高度平衡和颜色平衡属性。最后，我们讨论了左旋和右旋函数，用于维护红黑树的平衡属性。通过使用Rust来实现红黑树，我们可以利用该语言的性能和安全性功能来创建高效且可靠的数据结构。

什么是红黑树（RBT）？

红黑树是一种二叉搜索树，其中每个节点都被着色为红色或黑色。节点的颜色用于平衡树，以使从根节点到任何叶子节点的最长路径不超过从根节点到任何其他叶子节点的最短路径的两倍。此属性称为高度平衡属性。

红黑树是按照以下规则构建的

  1. 每个节点都是红色或黑色。
  2. 根节点始终为黑色。
  3. 每个叶子（NIL）节点都是黑色的。
  4. 如果一个节点为红色，则它的两个子节点必须为黑色。
  5. 从给定节点到其任何后代NIL节点的每条路径必须包含相同数量的黑色节点。

通过为每个节点分配颜色并限制连续的着色节点数量，RBT确保最长的分支永远不会超过最短分支的两倍大小，从而提供更稳定和有效的树。

以下是RBT的视觉表示：

photo_2023-02-22_21-58-43.jpg

当我们按顺序插入“1, 2, 3, 4, 5”时，以下是BST和RBT之间的比较，展示了为什么在插入排序元素时RBT远比BST有效：

photo_2023-02-22_22-32-17.jpg

很明显，当收到排序序列时，BST将变得非常低效。而RBT在此情况下可以更加平衡。

红黑树函数

红黑树支持以下函数：

  1. 插入：要将新节点插入红黑树中，我们首先创建一个具有给定键和值的新节点。然后，我们遍历树以找到插入新节点的正确位置。如果树为空，我们将新节点简单地作为树的根。如果树不为空，我们将新节点的键与当前正在检查的节点的键进行比较。如果新节点的键小于当前节点的键，则移动到当前节点的左子节点。如果新节点的键大于当前节点的键，则移动到当前节点的右子节点。我们重复这个过程，直到我们找到一个可以插入新节点的空位置。

  2. 删除：要从红黑树中删除节点，我们首先搜索具有给定键的节点。如果找不到节点，我们简单地返回而不做任何事情。如果找到节点，我们用其后继节点替换它，后继节点是节点右子树中键值最小的节点。然后，我们使用类似的过程从树中删除后继节点。

  3. 搜索：在红黑树中搜索节点类似于在二叉搜索树中搜索。我们从根节点开始，将给定键与当前节点的键进行比较。如果键相等，我们返回当前节点的值。如果给定键小于当前节点的键，我们移动到左子节点。如果给定键大于当前节点的键，则移动到右子节点。我们重复这个过程，直到我们找到具有给定键的节点或到达叶子节点。

Rust中的红黑树代码

要在Rust中实现红黑树，我们首先定义一个Node结构来表示树中的每个节点。Node结构包含节点的键，值，颜色和左右子节点字段。

1
2
3
4
5
6
7
8
9
10
11
enum Color {   // Defines an enum called 'Color'
    Red,      // Enum variant representing the color red
    Black,    // Enum variant representing the color black
}
struct Node<K, V> {    // Defines a struct called 'Node' with generic parameters 'K' and 'V'
    key: K,            // Field representing the key of the node
    value: V,          // Field representing the value of the node
    color: Color,      // Field representing the color of the node, which is of type 'Color' enum
    left: Option<Box<Node<K, V>>>,   // Field representing the left child node of the current node (if exists); it is an Option type and wrapped inside a box
    right: Option<Box<Node<K, V>>>,  // Field representing the right child node of the current node (if exists); it is an Option type and wrapped inside a box
}

我们还定义了一个RBTree结构来表示整个红黑树。RBTree结构包含一个根节点和用于插入，删除和搜索树中元素的方法。

1
2
3
pub struct RBTree<K: Ord, V> {   // Defines a struct called 'RBTree' with generic parameters 'K' and 'V'; 'K' must implement 'Ord' trait
    root: Option<Box<Node<K, V>>>,   // Field representing the root node of the red-black tree; it is an Option type and wrapped inside a box
}

左旋和右旋

这些函数执行给定节点及其子节点的左旋或右旋操作。rotate_left函数接收一个可变引用到一个节点，并将其右子节点向左旋转，而rotate_right函数接收一个可变引用到一个节点，并将其左子节点向右旋转。

在每次旋转期间，适当的指针被更新以反映树的新结构。每个节点的大小也被更新以反映旋转过程中发生的任何更改。

这些旋转函数由fix_after_insertion和fix_after_deletion方法用于维护红黑树的平衡属性。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
impl<K: Ord, V> RBTree<K, V> {

    // Defines a private method called 'rotate_left' that takes a mutable reference to a 'Node' struct
    fn rotate_left(mut node: &mut Node<K, V>) {
        // Takes ownership of the right child node of the given node
        let mut right = node.right.take().unwrap();
        
        // Moves the left child of the right child node to the right child of the given node
        node.right = right.left.take();
        
        // Sets the parent of the right child node as the given node and returns the right child node
        right.left = Some(std::mem::replace(&mut node, right));
        
        // Sets the parent of the left child node of the given node to the given node
        if let Some(ref mut left) = node.left {
            left.parent = Some(node);
        }
        
        // Sets the parent of the right child node of the given node to the given node
        if let Some(ref mut right) = node.right {
            right.parent = Some(node);
        }
        
        // Updates the size of the given node
        node.update_size();
    }

    // Defines a private method called 'rotate_right' that takes a mutable reference to a 'Node' struct
    fn rotate_right(mut node: &mut Node<K, V>) {
        // Takes ownership of the left child node of the given node
        let mut left = node.left.take().unwrap();
        
        // Moves the right child of the left child node to the left child of the given node
        node.left = left.right.take();
        
        // Sets the parent of the left child node as the given node and returns the left child node
        left.right = Some(std::mem::replace(&mut node, left));
        
        // Sets the parent of the left child node of the given node to the given node
        if let Some(ref mut left) = node.left {
            left.parent = Some(node);
        }
        
        // Sets the parent of the right child node of the given node to the given node
        if let Some(ref mut right) = node.right {
            right.parent = Some(node);
        }
        
        // Updates the size of the given node
        node.update_size();
    }
}

插入函数

插入函数遵循之前描述的插入函数相同的逻辑。我们创建一个具有给定键和值的新节点，并遍历树以找到插入新节点的正确位置。如果树为空，我们将新节点简单地作为树的根。如果树不为空，我们将新节点的键与当前正在检查的节点的键进行比较。如果新节点的键小于当前节点的键，则移动到当前节点的左子节点。如果新节点的键大于当前节点的键，则移动到当前节点的右子节点。我们重复这个过程，直到我们找到一个可以插入新节点的空位置。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
impl<K: Ord, V> RBTree<K, V> {

    // Defines a private method called 'fix_after_insertion' that takes a mutable reference to a 'Node' struct and an optional mutable reference to the root node
    fn fix_after_insertion(
        mut node: &mut Node<K, V>,
        root: Option<&mut Box<Node<K, V>>>,
    ) {
        // Sets the color of the inserted node as 'Red'
        node.color = Color::Red;
        
        // Checks if the node is not the root node and its parent's color is 'Red'
        while node != root.unwrap() && node.parent().unwrap().color == Color::Red {
            let parent = node.parent().unwrap();
            let grandparent = node.grandparent().unwrap();
            
            // Checks if the parent is the left child of its grandparent
            if parent == grandparent.left.as_ref().unwrap() {
                let uncle = grandparent.right.as_ref();
                // Checks if the uncle node exists and its color is 'Red'
                if uncle.is_some() && uncle.unwrap().color == Color::Red {
                    // Recolors the parent, uncle, and grandparent nodes, and sets the current node to its grandparent
                    parent.color = Color::Black;
                    uncle.unwrap().color = Color::Black;
                    grandparent.color = Color::Red;
                    node = grandparent;
                } else {
                    // Checks if the current node is the right child of its parent; if so, rotates it to the left around the parent node
                    if node == parent.right.as_ref().unwrap() {
                        node = parent;
                        RBTree::rotate_left(parent);
                    }
                    // Recolors the parent and grandparent nodes, and rotates the grandparent node to the right
                    parent.color = Color::Black;
                    grandparent.color = Color::Red;
                    RBTree::rotate_right(grandparent);
                }
            } else {
                let uncle = grandparent.left.as_ref();
                // Checks if the uncle node exists and its color is 'Red'
                if uncle.is_some() && uncle.unwrap().color == Color::Red {
                    // Recolors the parent, uncle, and grandparent nodes, and sets the current node to its grandparent
                    parent.color = Color::Black;
                    uncle.unwrap().color = Color::Black;
                    grandparent.color = Color::Red;
                    node = grandparent;
                } else {
                    // Checks if the current node is the left child of its parent; if so, rotates it to the right around the parent node
                    if node == parent.left.as_ref().unwrap() {
                        node = parent;
                        RBTree::rotate_right(parent);
                    }
                    // Recolors the parent and grandparent nodes, and rotates the grandparent node to the left
                    parent.color = Color::Black;
                    grandparent.color = Color::Red;
                    RBTree::rotate_left(grandparent);
                }
            }
        }
        root.unwrap().color = Color::Black;
    }

    // Defines a public method called 'insert' that takes a key and a value and adds it to the red-black tree
    pub fn insert(&mut self, key: K, value: V) {
        // Creates a new node with the given key, value, and color 'Red'
        let mut new_node = Box::new(Node {
            key,
            value,
            color: Color::Red,
            left: None,
            right: None,
        });

        // Checks if the root node exists
        if let Some(ref mut root) = self.root {
            let mut current = root.as_mut();
            // Traverses through the tree until it finds a suitable place to insert the new node
            loop {
                if new_node.key < current.key {
                    if let Some(ref mut left) = current.left {
                        current = left.as_mut();
                    } else {
                        current.left = Some(new_node);
                        break;
                    }
                } else if new_node.key > current.key {
                    if let Some(ref mut right) = current.right {
                        current = right.as_mut();
                    } else {
                        current.right = Some(new_node);
                        break;
                    }
                } else {
                    // If the key already exists in the tree, updates its corresponding value and exits the loop
                    current.value = new_node.value;
                    return;
                }
            }
            // Fixes the tree after insertion of the new node
            RBTree::fix_after_insertion(current, Some(&mut self.root));
        } else {
            // If the root node does not exist, sets the color of the new node as 'Black' and makes it the root node
            new_node.color = Color::Black;
            self.root = Some(new_node);
        }
    }
}

删除函数

删除函数也遵循之前描述的删除函数相同的逻辑。我们搜索具有给定键的节点，并用其后继节点替换它，后继节点是节点右子树中键值最小的节点。然后，我们使用类似的过程从树中删除后继节点。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// This function fixes the Red-Black Tree violations that may arise after a node deletion
fn fix_after_deletion(
    mut node: &mut Node<K, V>,
    root: Option<&mut Box<Node<K, V>>>,
) {
    while node != root.unwrap() && node.color == Color::Black {

        // Get the parent and sibling of the current node
        let parent = node.parent_mut().unwrap();
        let sibling = node.sibling().unwrap();

        if sibling.color == Color::Red {
            // Case 1: The sibling of the current node is red
            // Change the colors of the parent, sibling, and the child of the sibling
            sibling.color = Color::Black;
            parent.color = Color::Red;
            if node.is_left_child() {
                // Rotate left if the current node is the left child
                RBTree::rotate_left(parent);
                sibling.color = parent.right.as_ref().unwrap().color;
                parent.right.as_mut().unwrap().color = Color::Black;
                node = parent.left.as_mut().unwrap();
            } else {
                // Rotate right if the current node is the right child
                RBTree::rotate_right(parent);
                sibling.color = parent.left.as_ref().unwrap().color;
                parent.left.as_mut().unwrap().color = Color::Black;
                node = parent.right.as_mut().unwrap();
            }
        } else {
            if let (Some(left), Some(right)) = (sibling.left.as_ref(), sibling.right.as_ref()) {
                if left.color == Color::Black && right.color == Color::Black {
                    // Case 2: The sibling of the current node is black and both its children are black
                    // Change the colors of the sibling and move up to the parent
                    sibling.color = Color::Red;
                    node = parent;
                } else {
                    if node.is_left_child() && right.color == Color::Black {
                        // Case 3: The sibling of the current node is black and the right child is black
                        // Change the colors of the sibling and its left child, and rotate right
                        sibling.color = Color::Red;
                        left.color = Color::Black;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    } else if node.is_right_child() && left.color == Color::Black {
                        // Case 3: The sibling of the current node is black and the left child is black
                        // Change the colors of the sibling and its right child, and rotate left
                        sibling.color = Color::Red;
                        right.color = Color::Black;
                        RBTree::rotate_left(sibling);
                        sibling = parent.left.as_mut().unwrap();
                    }

                    // Case 4: The sibling of the current node is black and has a red child
                    // Change the colors of the parent, sibling, and the child of the sibling
                    sibling.color = parent.color;
                    parent.color = Color::Black;
                    if node.is_left_child() {
                        right.color = Color::Black;
                        RBTree::rotate_left(parent);
                    } else {
                        left.color = Color::Black;
                        RBTree::rotate_right(parent);
                    }
                    break;
                }
            } else {
                // Case 2: The sibling of the current node is black and has no children
                if let Some(left) = sibling.left.as_ref() {
                    if node.is_left_child() {
                        // Change the color of the sibling's left child and rotate right
                        left.color = Color::Black;
                        sibling.color = parent.color;
                        RBTree::rotate_right(parent);
                    } else {
                        // Change the color of the sibling and its left child, and rotate right
                        left.color = Color::Red;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    }
                } else {
                    if node.is_left_child() {
                        // Change the color of the sibling and its left child, and rotate right
                        sibling.color = Color::Red;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    } else {
                        // Change the colors of the parent and sibling, and rotate left
                        sibling.color = parent.color;
                        parent.color = Color::Black;
                        RBTree::rotate_left(parent);
                        break;
                    }
                }
            }
        }
    }

    // Set the color of the current node to black
    node.color = Color::Black;
}
// This function deletes a node from the Red-Black Tree according to the given key
pub fn delete(&mut self, key: &K) -> Option<V> {
    let mut current = self.root.as_mut();

    while let Some(node) = current {
        if key < &node.key {
            current = node.left.as_mut();
        } else if key > &node.key {
            current = node.right.as_mut();
        } else {
            if node.left.is_some() && node.right.is_some() {
                // If the node has two children, replace it with its successor and delete the successor
                let successor = node.right.as_mut().unwrap().min_node();
                node.key = successor.key;
                node.value = std::mem::replace(&mut successor.value, Default::default());
                current = Some(&mut *successor);
            } else {
                // If the node has one child or no children, remove it and fix any Red-Black Tree violations
                let child = if node.left.is_some() {
                    node.left.take()
                } else {
                    node.right.take()
                };

                if node.color == Color::Black {
                    if let Some(ref c) = child {
                        RBTree::fix_after_deletion(c.as_mut(), Some(&mut self.root));
                    } else {
                        RBTree::fix_after_deletion(node, Some(&mut self.root));
                    }
                }

                return Some(node.value);
            }
        }
    }

    None
}

搜索函数

搜索函数与之前描述的搜索函数类似。我们从根节点开始，将给定键与当前节点的键进行比较。如果键相等，我们返回当前节点的值。如果给定键小于当前节点的键，我们移动到左子节点。如果给定键大于当前节点的键，则移动到右子节点。我们重复这个过程，直到我们找到具有给定键的节点或到达叶子节点。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
impl<K: Ord, V> RBTree<K, V> {
    pub fn search(&self, key: &K) -> Option<&V> {
        let mut current = self.root.as_ref(); // Start search from the root of the tree.
        
        while let Some(node) = current { // Traverse the tree until the key is found or a leaf node is reached.
            if key == &node.key { // If the key is found, return the corresponding value.
                return Some(&node.value);
            } else if key < &node.key { // If the key is less than the current node's key, search the left subtree.
                current = node.left.as_ref();
            } else { // If the key is greater than the current node's key, search the right subtree.
                current = node.right.as_ref();
            }
        }

        None // If the key is not found, return None.
    }
}

结论

红黑树是一种自平衡的二叉搜索树，用于按排序顺序存储和检索元素，所有主要操作的时间复杂度为O(log n)。它们旨在具有内存效率和平衡性，使其比其他平衡树如AVL树更有效率。Rust是一种现代编程语言，旨在提高性能和安全性，因此它是实现数据结构如红黑树的绝佳选择。

Posted on Edited on In

中文版链接

Abstract

In computer science, Red-Black Trees are a type of self-balancing binary search tree that are used to store and retrieve elements in sorted order with O(log n) time complexity for all major operations. They are designed to be memory-efficient and balanced, making them more efficient than other balanced trees like AVL trees. In this article, we have discussed the implementation of Red-Black Trees in Rust, a modern programming language designed for performance and safety. We have covered the structure of the tree, as well as the functions for insertion, deletion, and searching. Additionally, we have covered the fix functions that are used to restore the height balance and color balance properties of the tree after each insertion and deletion. Finally, we have discussed the left and right rotation functions that are used to maintain the balance properties of the Red-Black Tree. By using Rust to implement Red-Black Trees, we can take advantage of the language’s performance and safety features to create an efficient and reliable data structure.

What is Red Black Tree (RBT)?

Red-Black Trees are binary search trees where each node is colored either red or black. The color of the node is used to balance the tree so that the longest path from the root to any leaf node is no more than twice as long as the shortest path from the root to any other leaf node. This property is known as the height balance property.

Red-Black Trees are constructed with the following rules

  1. Each node is either red or black.
  2. The root node is always black.
  3. Every leaf (NIL) node is black.
  4. If a node is red, then both its children must be black.
  5. Every path from a given node to any of its descendant NIL nodes must contain the same number of black nodes.

By assigning each node with a color and by limiting the number of consecutive colored nodes, RBT ensures that the longest branches will never exceed twice the size of the shortest branches, thus providing a more stable and effective tree. For reference, here’s a visual representation of an RBT:

photo_2023-02-22_21-58-43.jpg

Here’s a comparison between BST and RBT when we insert “1, 2, 3, 4, 5” in order, demonstrating why RBT is far more efficient than BST when sorted elements are inserted:

photo_2023-02-22_22-32-17.jpg

Obviously, BST will become extremely low effective when receiving a sorted sequence. And RBT can become more balanced under this situation.

Red-Black Tree Functions

A Red-Black Tree supports the following functions:

  1. Insertion: To insert a new node into the Red-Black Tree, we first create a new node with the given key and value. We then traverse the tree to find the correct location to insert the new node. If the tree is empty, we simply make the new node the root of the tree. If the tree is not empty, we compare the new node’s key to the key of the current node we are examining. If the new node’s key is less than the current node’s key, we move to the left child of the current node. If the new node’s key is greater than the current node’s key, we move to the right child of the current node. We repeat this process until we find an empty location in the tree where we can insert the new node.

  2. Deletion: To delete a node from the Red-Black Tree, we first search for the node with the given key. If the node is not found, we simply return without doing anything. If the node is found, we replace it with its successor, which is the node with the smallest key in the node’s right subtree. We then delete the successor node from the tree using a similar process.

  3. Searching: Searching for a node in the Red-Black Tree is similar to searching in a binary search tree. We start at the root node and compare the given key with the key of the current node. If the keys are equal, we return the value of the current node. If the given key is less than the key of the current node, we move to the left child. If the given key is greater than the key of the current node, we move to the right child. We repeat this process until we find the node with the given key or reach a leaf node.

Rust Code for Red-Black Trees

To implement Red-Black Trees in Rust, we first define a Node struct to represent each node in the tree. The Node struct contains fields for the node’s key, value, color, and left and right children.

1
2
3
4
5
6
7
8
9
10
11
enum Color {   // Defines an enum called 'Color'
    Red,      // Enum variant representing the color red
    Black,    // Enum variant representing the color black
}
struct Node<K, V> {    // Defines a struct called 'Node' with generic parameters 'K' and 'V'
    key: K,            // Field representing the key of the node
    value: V,          // Field representing the value of the node
    color: Color,      // Field representing the color of the node, which is of type 'Color' enum
    left: Option<Box<Node<K, V>>>,   // Field representing the left child node of the current node (if exists); it is an Option type and wrapped inside a box
    right: Option<Box<Node<K, V>>>,  // Field representing the right child node of the current node (if exists); it is an Option type and wrapped inside a box
}

We also define an RBTree struct to represent the Red-Black Tree as a whole. The RBTree struct contains a root node and methods to insert, delete, and search for elements in the tree.

1
2
3
pub struct RBTree<K: Ord, V> {   // Defines a struct called 'RBTree' with generic parameters 'K' and 'V'; 'K' must implement 'Ord' trait
    root: Option<Box<Node<K, V>>>,   // Field representing the root node of the red-black tree; it is an Option type and wrapped inside a box
}

Left and Right Rotate

These functions perform a left or right rotation of a given node and its children. The rotate_left function takes in a mutable reference to a node and rotates its right child to the left, while the rotate_right function takes in a mutable reference to a node and rotates its left child to the right.

During each rotation, the appropriate pointers are updated to reflect the new structure of the tree. The size of each node is also updated to reflect any changes that occurred during the rotation.

These rotation functions are used by the fix_after_insertion and fix_after_deletion methods to maintain the balance properties of the Red-Black Tree.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
impl<K: Ord, V> RBTree<K, V> {

    // Defines a private method called 'rotate_left' that takes a mutable reference to a 'Node' struct
    fn rotate_left(mut node: &mut Node<K, V>) {
        // Takes ownership of the right child node of the given node
        let mut right = node.right.take().unwrap();
        
        // Moves the left child of the right child node to the right child of the given node
        node.right = right.left.take();
        
        // Sets the parent of the right child node as the given node and returns the right child node
        right.left = Some(std::mem::replace(&mut node, right));
        
        // Sets the parent of the left child node of the given node to the given node
        if let Some(ref mut left) = node.left {
            left.parent = Some(node);
        }
        
        // Sets the parent of the right child node of the given node to the given node
        if let Some(ref mut right) = node.right {
            right.parent = Some(node);
        }
        
        // Updates the size of the given node
        node.update_size();
    }

    // Defines a private method called 'rotate_right' that takes a mutable reference to a 'Node' struct
    fn rotate_right(mut node: &mut Node<K, V>) {
        // Takes ownership of the left child node of the given node
        let mut left = node.left.take().unwrap();
        
        // Moves the right child of the left child node to the left child of the given node
        node.left = left.right.take();
        
        // Sets the parent of the left child node as the given node and returns the left child node
        left.right = Some(std::mem::replace(&mut node, left));
        
        // Sets the parent of the left child node of the given node to the given node
        if let Some(ref mut left) = node.left {
            left.parent = Some(node);
        }
        
        // Sets the parent of the right child node of the given node to the given node
        if let Some(ref mut right) = node.right {
            right.parent = Some(node);
        }
        
        // Updates the size of the given node
        node.update_size();
    }
}

Insertion Function in Rust

The insert function in Rust follows the same logic as the insertion function described earlier. We create a new node with the given key and value and traverse the tree to find the correct location to insert the new node. If the tree is empty, we simply make the new node the root of the tree. If the tree is not empty, we compare the new node’s key to the key of the current node we are examining. If the new node’s key is less than the current node’s key, we move to the left child of the current node. If the new node’s key is greater than the current node’s key, we move to the right child of the current node. We repeat this process until we find an empty location in the tree where we can insert the new node.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
impl<K: Ord, V> RBTree<K, V> {

    // Defines a private method called 'fix_after_insertion' that takes a mutable reference to a 'Node' struct and an optional mutable reference to the root node
    fn fix_after_insertion(
        mut node: &mut Node<K, V>,
        root: Option<&mut Box<Node<K, V>>>,
    ) {
        // Sets the color of the inserted node as 'Red'
        node.color = Color::Red;
        
        // Checks if the node is not the root node and its parent's color is 'Red'
        while node != root.unwrap() && node.parent().unwrap().color == Color::Red {
            let parent = node.parent().unwrap();
            let grandparent = node.grandparent().unwrap();
            
            // Checks if the parent is the left child of its grandparent
            if parent == grandparent.left.as_ref().unwrap() {
                let uncle = grandparent.right.as_ref();
                // Checks if the uncle node exists and its color is 'Red'
                if uncle.is_some() && uncle.unwrap().color == Color::Red {
                    // Recolors the parent, uncle, and grandparent nodes, and sets the current node to its grandparent
                    parent.color = Color::Black;
                    uncle.unwrap().color = Color::Black;
                    grandparent.color = Color::Red;
                    node = grandparent;
                } else {
                    // Checks if the current node is the right child of its parent; if so, rotates it to the left around the parent node
                    if node == parent.right.as_ref().unwrap() {
                        node = parent;
                        RBTree::rotate_left(parent);
                    }
                    // Recolors the parent and grandparent nodes, and rotates the grandparent node to the right
                    parent.color = Color::Black;
                    grandparent.color = Color::Red;
                    RBTree::rotate_right(grandparent);
                }
            } else {
                let uncle = grandparent.left.as_ref();
                // Checks if the uncle node exists and its color is 'Red'
                if uncle.is_some() && uncle.unwrap().color == Color::Red {
                    // Recolors the parent, uncle, and grandparent nodes, and sets the current node to its grandparent
                    parent.color = Color::Black;
                    uncle.unwrap().color = Color::Black;
                    grandparent.color = Color::Red;
                    node = grandparent;
                } else {
                    // Checks if the current node is the left child of its parent; if so, rotates it to the right around the parent node
                    if node == parent.left.as_ref().unwrap() {
                        node = parent;
                        RBTree::rotate_right(parent);
                    }
                    // Recolors the parent and grandparent nodes, and rotates the grandparent node to the left
                    parent.color = Color::Black;
                    grandparent.color = Color::Red;
                    RBTree::rotate_left(grandparent);
                }
            }
        }
        root.unwrap().color = Color::Black;
    }

    // Defines a public method called 'insert' that takes a key and a value and adds it to the red-black tree
    pub fn insert(&mut self, key: K, value: V) {
        // Creates a new node with the given key, value, and color 'Red'
        let mut new_node = Box::new(Node {
            key,
            value,
            color: Color::Red,
            left: None,
            right: None,
        });

        // Checks if the root node exists
        if let Some(ref mut root) = self.root {
            let mut current = root.as_mut();
            // Traverses through the tree until it finds a suitable place to insert the new node
            loop {
                if new_node.key < current.key {
                    if let Some(ref mut left) = current.left {
                        current = left.as_mut();
                    } else {
                        current.left = Some(new_node);
                        break;
                    }
                } else if new_node.key > current.key {
                    if let Some(ref mut right) = current.right {
                        current = right.as_mut();
                    } else {
                        current.right = Some(new_node);
                        break;
                    }
                } else {
                    // If the key already exists in the tree, updates its corresponding value and exits the loop
                    current.value = new_node.value;
                    return;
                }
            }
            // Fixes the tree after insertion of the new node
            RBTree::fix_after_insertion(current, Some(&mut self.root));
        } else {
            // If the root node does not exist, sets the color of the new node as 'Black' and makes it the root node
            new_node.color = Color::Black;
            self.root = Some(new_node);
        }
    }
}

Deletion Function in Rust

The delete function in Rust also follows the same logic as the deletion function described earlier. We search for the node with the given key and replace it with its successor, which is the node with the smallest key in the node’s right subtree. We then delete the successor node from the tree using a similar process.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// This function fixes the Red-Black Tree violations that may arise after a node deletion
fn fix_after_deletion(
    mut node: &mut Node<K, V>,
    root: Option<&mut Box<Node<K, V>>>,
) {
    while node != root.unwrap() && node.color == Color::Black {

        // Get the parent and sibling of the current node
        let parent = node.parent_mut().unwrap();
        let sibling = node.sibling().unwrap();

        if sibling.color == Color::Red {
            // Case 1: The sibling of the current node is red
            // Change the colors of the parent, sibling, and the child of the sibling
            sibling.color = Color::Black;
            parent.color = Color::Red;
            if node.is_left_child() {
                // Rotate left if the current node is the left child
                RBTree::rotate_left(parent);
                sibling.color = parent.right.as_ref().unwrap().color;
                parent.right.as_mut().unwrap().color = Color::Black;
                node = parent.left.as_mut().unwrap();
            } else {
                // Rotate right if the current node is the right child
                RBTree::rotate_right(parent);
                sibling.color = parent.left.as_ref().unwrap().color;
                parent.left.as_mut().unwrap().color = Color::Black;
                node = parent.right.as_mut().unwrap();
            }
        } else {
            if let (Some(left), Some(right)) = (sibling.left.as_ref(), sibling.right.as_ref()) {
                if left.color == Color::Black && right.color == Color::Black {
                    // Case 2: The sibling of the current node is black and both its children are black
                    // Change the colors of the sibling and move up to the parent
                    sibling.color = Color::Red;
                    node = parent;
                } else {
                    if node.is_left_child() && right.color == Color::Black {
                        // Case 3: The sibling of the current node is black and the right child is black
                        // Change the colors of the sibling and its left child, and rotate right
                        sibling.color = Color::Red;
                        left.color = Color::Black;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    } else if node.is_right_child() && left.color == Color::Black {
                        // Case 3: The sibling of the current node is black and the left child is black
                        // Change the colors of the sibling and its right child, and rotate left
                        sibling.color = Color::Red;
                        right.color = Color::Black;
                        RBTree::rotate_left(sibling);
                        sibling = parent.left.as_mut().unwrap();
                    }

                    // Case 4: The sibling of the current node is black and has a red child
                    // Change the colors of the parent, sibling, and the child of the sibling
                    sibling.color = parent.color;
                    parent.color = Color::Black;
                    if node.is_left_child() {
                        right.color = Color::Black;
                        RBTree::rotate_left(parent);
                    } else {
                        left.color = Color::Black;
                        RBTree::rotate_right(parent);
                    }
                    break;
                }
            } else {
                // Case 2: The sibling of the current node is black and has no children
                if let Some(left) = sibling.left.as_ref() {
                    if node.is_left_child() {
                        // Change the color of the sibling's left child and rotate right
                        left.color = Color::Black;
                        sibling.color = parent.color;
                        RBTree::rotate_right(parent);
                    } else {
                        // Change the color of the sibling and its left child, and rotate right
                        left.color = Color::Red;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    }
                } else {
                    if node.is_left_child() {
                        // Change the color of the sibling and its left child, and rotate right
                        sibling.color = Color::Red;
                        RBTree::rotate_right(sibling);
                        sibling = parent.right.as_mut().unwrap();
                    } else {
                        // Change the colors of the parent and sibling, and rotate left
                        sibling.color = parent.color;
                        parent.color = Color::Black;
                        RBTree::rotate_left(parent);
                        break;
                    }
                }
            }
        }
    }

    // Set the color of the current node to black
    node.color = Color::Black;
}
// This function deletes a node from the Red-Black Tree according to the given key
pub fn delete(&mut self, key: &K) -> Option<V> {
    let mut current = self.root.as_mut();

    while let Some(node) = current {
        if key < &node.key {
            current = node.left.as_mut();
        } else if key > &node.key {
            current = node.right.as_mut();
        } else {
            if node.left.is_some() && node.right.is_some() {
                // If the node has two children, replace it with its successor and delete the successor
                let successor = node.right.as_mut().unwrap().min_node();
                node.key = successor.key;
                node.value = std::mem::replace(&mut successor.value, Default::default());
                current = Some(&mut *successor);
            } else {
                // If the node has one child or no children, remove it and fix any Red-Black Tree violations
                let child = if node.left.is_some() {
                    node.left.take()
                } else {
                    node.right.take()
                };

                if node.color == Color::Black {
                    if let Some(ref c) = child {
                        RBTree::fix_after_deletion(c.as_mut(), Some(&mut self.root));
                    } else {
                        RBTree::fix_after_deletion(node, Some(&mut self.root));
                    }
                }

                return Some(node.value);
            }
        }
    }

    None
}

Searching Function in Rust

The search function in Rust is also similar to the search function described earlier. We start at the root node and compare the given key with the key of the current node. If the keys are equal, we return the value of the current node. If the given key is less than the key of the current node, we move to the left child. If the given key is greater than the key of the current node, we move to the right child. We repeat this process until we find the node with the given key or reach a leaf node.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
impl<K: Ord, V> RBTree<K, V> {
    pub fn search(&self, key: &K) -> Option<&V> {
        let mut current = self.root.as_ref(); // Start search from the root of the tree.
        
        while let Some(node) = current { // Traverse the tree until the key is found or a leaf node is reached.
            if key == &node.key { // If the key is found, return the corresponding value.
                return Some(&node.value);
            } else if key < &node.key { // If the key is less than the current node's key, search the left subtree.
                current = node.left.as_ref();
            } else { // If the key is greater than the current node's key, search the right subtree.
                current = node.right.as_ref();
            }
        }

        None // If the key is not found, return None.
    }
}

Conclusion

Red-Black Trees are a type of self-balancing binary search tree that are used to store and retrieve elements in sorted order with O(log n) time complexity for all major operations. They are designed to be memory efficient and balanced, making them more efficient than other balanced trees like AVL trees. Rust is a modern programming language that is designed for performance and safety, making it a great choice for implementing data structures like the Red-Black Tree.

0%