When dealing with recursive data structures, you can often expect some recursive code to work on them. This case is not an exception either, just the recursive parts will need some dirtiness.
Let's use the 3-element BSTs for designing, labeled with insertion-order:
123 132 213,231 312 321
1 1 2 3 3
\ \ / \ / /
2 3 1 3 1 2
\ / \ /
3 2 2 1
Finding the largest element is easy:
- just go to the right as long as you can
You will bump into 3, no matter what level it is.
Finding the second largest element is the revealing part:
- going to the right as long as it's possible still seems to be a good start
- then we look at where we are:
- if it's a leaf node, return to the parent, and that will be the one (123,213,231 cases)
- if there's a left-child, check that one, but as the 312 case in particular shows, "checking" the left-child actually means step 1, so again go to the right as long as it's possible.
The recursion is somewhat found, and these really are the steps we need for larger cases too. It's also somewhat seen that when we are going to the right, the "nth-ness" of the next number we will check doesn't change. It starts changing only when we are stepping to the left (132,312,321), or returning to a previous level (123,213,231).
The dark part is that we have to track this counter somehow. We found the answer when it reaches 0 (so starting this algorithm with n=0 finds the largest element), and after that (when n goes negative) it will just return the value it got from recursion.
First here is a JavaScript PoC, using a bit dirty hacks, like if a member variable doesn't exist at all yet, it still can be checked (the if(this.left) things), and the counter is a one-element array (so it can be modified across the recursive calls), the method is called as root.nthback([i]), where the [i] is an array literal. Also, the method doesn't bother returning anything when the element doesn't exist, that produces the two undefineds at the end of the test output. These shortcuts will be addressed in the Java variant at the end of this post.
The example input was just taken from the other answer, on top of their availability, they have some repeats too.
const init = [148, 65, 18, 168, 8, 148, 194, 186, 114, 22, 102, 51, 123, 169, 68, 118, 37, 18, 26, 18];
class Node {
nthback(n) {
if (this.right) {
let res = this.right.nthback(n);
if (n[0] < 0)
return res;
}
if (n[0]-- === 0)
return this.value;
if (this.left) {
let res = this.left.nthback(n);
if (n[0] < 0)
return res;
}
}
constructor() {
this.value = NaN;
}
add(value) {
if (isNaN(this.value)) {
this.value = value;
} else if (value < this.value) {
if (!this.left)
this.left = new Node;
this.left.add(value);
} else {
if (!this.right)
this.right = new Node;
this.right.add(value);
}
}
walk() {
let result = "";
if (this.left)
result = this.left.walk() + ",";
result += this.value;
if (this.right)
result += "," + this.right.walk();
return result;
}
}
const root = new Node;
for (const value of init)
root.add(value);
console.log(root.walk());
for (let i = 0; i < 22; i++)
console.log(root.nthback([i]));
So the actual magic is quite short and also symmetric:
nthback(n) {
if (this.right) { // 1
const res = this.right.nthback(n); // 2
if (n[0] < 0) // 3
return res;
}
if (n[0]-- === 0) // 4
return this.value;
if (this.left) { // 5
const res = this.left.nthback(n);
if (n[0] < 0)
return res;
}
}
If there is something on the right (1), it has to be checked (2), and if the counter is negative afterwards (3), the result we got back is the actual result of the entire call, so we pass it back.
If we are still in the method, (4) is where we check if the counter is exactly 0, because then this node has the actual result, which we can return. It's worth to remember that the n[0]-- part decrements the counter regardless of the outcome of the comparison. So if n[0] was 0 initially, it will become -1 and we return this.value;. If it was something else, it just gets decremented.
(5) does the (1)-(2)-(3) part for the left branch. And the hidden JavaScript thing is that we don't have to return anything. But the (4)-(5) parts will change when using your complete structure anyway.
Adding subtree-size allows early return if the incoming counter is simply larger than the size of the entire subtree: we just decrement the counter by the size, and return, the element is somewhere else. And this also means that when we don't return, the result is in the subtree, so we check the possible right branch, then ourselves, and if we are still inside the method, we don't even have to check if we have a left branch, because we do have it for sure, and it does contain the result, also for sure. So (5) will simply become a direct return this.left.nthback(n);. Which is quite a simplification.
Tracking multiplicity affects (4): instead of checking for 0, we will have to check if the counter is less than the multiplicity, and also, instead of decrementing the counter by 1, we have to subtract the actual multiplicity from it.
const init = [148, 65, 18, 168, 8, 148, 194, 186, 114, 22, 102, 51, 123, 169, 68, 118, 37, 18, 26, 18];
class Node {
nthback(n) {
if (this.size <= n[0]) {
n[0] -= this.size;
return 0;
}
if (this.right) {
let res = this.right.nthback(n);
if (n[0] < 0)
return res;
}
if (n[0] < this.count) {
n[0] -= this.count;
return this.value;
}
n[0] -= this.count;
return this.left.nthback(n);
}
constructor() {
this.value = NaN;
this.size = 0;
}
add(value) {
this.size++;
if (isNaN(this.value)) {
this.value = value;
this.count = 1;
} else if (value === this.value) {
this.count++;
} else if (value < this.value) {
if (!this.left)
this.left = new Node;
this.left.add(value);
} else {
if (!this.right)
this.right = new Node;
this.right.add(value);
}
}
walk() {
let result = "";
if (this.left)
result = this.left.walk() + ",";
result += this.value;
if (this.count > 1)
result += "x" + this.count;
result += "(" + this.size + ")";
if (this.right)
result += "," + this.right.walk();
return result;
}
}
const root = new Node;
for (const value of init)
root.add(value);
console.log(root.walk());
for (let i = 0; i < 22; i++)
console.log(root.nthback([i]));
So the final JavaScript variant could look like this:
nthback(n) {
if (this.size <= n[0]) { // 1
n[0] -= this.size;
return 0;
}
if (this.right) { // 2
let res = this.right.nthback(n);
if (n[0] < 0)
return res;
}
if (n[0] < this.count) { // 3
n[0] -= this.count; // 4
return this.value;
}
n[0] -= this.count; // 4
return this.left.nthback(n); // 5
}
Subtree-skipping happens in (1), by comparing subtree-size, and the target count we can immediately tell if the result is in this subtree or somewhere else. While JavaScript would allow a simple return; here, a 0 is produced instead, as it seems to be desired in the question.
The next step (2) is unchanged from the previous variant, looks exactly same, does exactly same.
(3) had to be taken apart. While the post-decrement operator was helpful previously, there's no post--= operator, so it happens for both outcomes of the comparison (4). Also, we don't compare === 0 as before, but we compare < this.count instead. By the way, we could have done that in the previous case too if we really wanted to, there this.count was always 1, so < 1 could have done the thing (the counter is never negative at this point, that can only happen in the left/right returns). In fact if we really-really wanted to, we could do the subtraction prior to the comparison, and instead of the current 0<=n fact and n<count check, we could shift the comparison downwards, knowing that now -count<=n and checking for n<0.
(5) became the one-liner as promised. If we reach this point, the result is unconditionally inside the left-branch, which unconditionally exists.
Making it Java
The simplest part is the array-trick: there could be a public method to call, accepting an actual number, and then it could wrap it into an array, and call a private method doing the actual job. Also changed the variable names to the ones you have:
public int nthback(int n) {
return nthback(new int[] { n });
}
private int nthback(int[] n) {
if (treeSize <= n[0]) {
n[0] -= treeSize;
return 0;
}
if (rightChild != null) {
int res = rightChild.nthback(n);
if (n[0] < 0)
return res;
}
if (n[0] < dataCount) {
n[0] -= dataCount;
return data;
}
n[0] -= dataCount;
return leftChild.nthback(n);
}
As you have private members, these methods have to reside in the same source file, and then they could just reside directly inside class BSTNode anyway.