Maximum Capacity Within Budget solution and explanation

class Solution:
    def maxCapacity(self, costs: List[int], capacity: List[int], budget: int) -> int:
        arr = []
        for a, b in zip(costs, capacity):
            if a < budget:
                arr.append((a, b))
        if not arr:
            return 0
        arr.sort()
        remain = SortedList()
        for i, (_, b) in enumerate(arr):
            remain.add((b, i))
        i, j = 0, len(arr) - 1
        ans = remain[-1][0]
        while i < j:
            remain.discard((arr[i][1], i))
            while i < j and arr[i][0] + arr[j][0] >= budget:
                remain.discard((arr[j][1], j))
                j -= 1
            if remain:
                ans = max(ans, arr[i][1] + remain[-1][0])
            i += 1
        return ans

class Solution {
    public int maxCapacity(int[] costs, int[] capacity, int budget) {
        List<int[]> arr = new ArrayList<>();
        for (int k = 0; k < costs.length; k++) {
            int a = costs[k], b = capacity[k];
            if (a < budget) {
                arr.add(new int[] {a, b});
            }
        }
        if (arr.isEmpty()) {
            return 0;
        }
        arr.sort(Comparator.comparingInt(o -> o[0]));
        TreeSet<int[]> remain = new TreeSet<>((x, y) -> {
            if (x[0] != y[0]) {
                return x[0] - y[0];
            }
            return x[1] - y[1];
        });
        for (int i = 0; i < arr.size(); i++) {
            remain.add(new int[] {arr.get(i)[1], i});
        }
        int i = 0, j = arr.size() - 1;
        int ans = remain.last()[0];
        while (i < j) {
            remain.remove(new int[] {arr.get(i)[1], i});
            while (i < j && arr.get(i)[0] + arr.get(j)[0] >= budget) {
                remain.remove(new int[] {arr.get(j)[1], j});
                j--;
            }
            if (!remain.isEmpty()) {
                ans = Math.max(ans, arr.get(i)[1] + remain.last()[0]);
            }
            i++;
        }
        return ans;
    }
}

class Solution {
public:
    int maxCapacity(vector<int>& costs, vector<int>& capacity, int budget) {
        vector<pair<int, int>> arr;
        for (int k = 0; k < costs.size(); k++) {
            int a = costs[k], b = capacity[k];
            if (a < budget) {
                arr.emplace_back(a, b);
            }
        }
        if (arr.empty()) {
            return 0;
        }
        sort(arr.begin(), arr.end());
        multiset<pair<int, int>> remain;
        for (int i = 0; i < arr.size(); i++) {
            remain.insert({arr[i].second, i});
        }
        int i = 0, j = arr.size() - 1;
        int ans = prev(remain.end())->first;
        while (i < j) {
            remain.erase(remain.find({arr[i].second, i}));
            while (i < j && arr[i].first + arr[j].first >= budget) {
                remain.erase(remain.find({arr[j].second, j}));
                j--;
            }
            if (!remain.empty()) {
                ans = max(ans, arr[i].second + prev(remain.end())->first);
            }
            i++;
        }
        return ans;
    }
};

type Node struct {
	b int
	i int
}

type MaxHeap []Node

func (h MaxHeap) Len() int { return len(h) }
func (h MaxHeap) Less(i, j int) bool {
	if h[i].b != h[j].b {
		return h[i].b > h[j].b
	}
	return h[i].i > h[j].i
}
func (h MaxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(Node))
}
func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func maxCapacity(costs []int, capacity []int, budget int) int {
	arr := make([][2]int, 0)
	for k := 0; k < len(costs); k++ {
		a, b := costs[k], capacity[k]
		if a < budget {
			arr = append(arr, [2]int{a, b})
		}
	}
	if len(arr) == 0 {
		return 0
	}
	sort.Slice(arr, func(i, j int) bool {
		if arr[i][0] != arr[j][0] {
			return arr[i][0] < arr[j][0]
		}
		return arr[i][1] < arr[j][1]
	})
	alive := make([]bool, len(arr))
	h := &MaxHeap{}
	for i := 0; i < len(arr); i++ {
		alive[i] = true
		heap.Push(h, Node{arr[i][1], i})
	}
	i, j := 0, len(arr)-1
	for h.Len() > 0 && !alive[(*h)[0].i] {
		heap.Pop(h)
	}
	ans := (*h)[0].b
	for i < j {
		alive[i] = false
		for i < j && arr[i][0]+arr[j][0] >= budget {
			alive[j] = false
			j--
		}
		for h.Len() > 0 && !alive[(*h)[0].i] {
			heap.Pop(h)
		}
		if h.Len() > 0 {
			if arr[i][1]+(*h)[0].b > ans {
				ans = arr[i][1] + (*h)[0].b
			}
		}
		i++
	}
	return ans
}