如何在Java中实现一个算法,该算法在O(1)时间内返回堆栈上前k个元素的总和

pinkon5k  于 12个月前  发布在  Java
关注(0)|答案(2)|浏览(80)

我正在尝试创建一个支持ksum()方法的堆栈的ArrayList实现。ksum()需要在O(1)时间内返回堆栈上前k个元素的总和。我应该使用什么算法来实现这一点?
目前我正在循环遍历前k个元素,并将它们的和添加到一个变量中,但是这需要O(n)时间。什么是更好的方法来获得顶部k元素的总和在一个堆栈在常数时间?

lymnna71

lymnna711#

除了实际值之外,还添加第二个变量来跟踪当前部分和。这样,得到顶部k元素的和就简单地由下式给出:

element[last].partial - element[last-k].partial
ig9co6j1

ig9co6j12#

我假设你在这里使用整数。您可以将位置i的堆栈的总和存储在位置i + 1的不同sum堆栈中(sum堆栈最初包含值0,因为它是加法的中性元素)。
当将一个int元素压入堆栈时,你可以通过添加新的总和来更新sum堆栈(只需添加你的int值和sum堆栈的当前顶部值)。
计算顶部k元素的总和可以通过从sum堆栈中的顶部总和值中减去k + 1总和值来实现。
下面是一个示例:

public class StackNSum {

    private final Stack<Integer> ints = new Stack<>();
    private final Stack<Integer> sums = new Stack<>();

    public StackNSum() {
        sums.push(0);
    }

    public void push(final int number) {
        sums.push(number + sums.peek());
        ints.push(number);
    }

    public int ksum(final int k) {
        // ints.size() - k < 0 need to be handled
        return sums.peek() - sums.get(ints.size() - k);
    }
}

相关问题