LeetCode/go/ipo.go

package main

import (
	"cmp"
	"slices"

	pq "github.com/emirpasic/gods/v2/queues/priorityqueue"
)

type Project struct {
	capital int
	profit  int
}

func byCapital(a, b Project) int {
	return cmp.Compare(a.capital, b.capital)
}

func revCmp[T cmp.Ordered](x, y T) int {
	return -cmp.Compare(x, y)
}

func findMaximizedCapital(k int, w int, profits []int, capital []int) int {
	// construct the list of projects
	projects := make([]Project, len(profits))
	for i, _ := range profits {
		projects[i] = Project{capital: capital[i], profit: profits[i]}
	}
	slices.SortFunc(projects, byCapital)

	queue := pq.NewWith[int](revCmp)

	i := 0
	for range k {
		// add the affordable projects to the queue
		for i < len(projects) && projects[i].capital <= w {
			queue.Enqueue(projects[i].profit)
			i++
		}

		// invest into the currently most profitable project
		profit, ok := queue.Dequeue()
		if !ok {
			// ran out of the available projects
			break
		}

		// collect the profit
		w += profit
	}

	return w
}