在一个数组中高效地找到子数组的算术平均值

Question

我正在尝试找到计算数组子数组算术平均值的方法。

问题可以归结为：给定数组X和整数S，有多少个连续的X片段的算术平均值等于S？例如，对于 X=[5,3,6,2] 和 S=4，结果为3。[5,3]，[6,2]和[5,3,6,2]的平均值都为4。

X可能最多有100,000个元素。每个X值都是{-1,000,000,000，+1,000,000,000}范围内的整数。S也是如此。我们不会四舍五入计算出的算术平均值。

下面是我的Java代码，适用于小数据集但效率不高，O(n^2)。

public static int returnSubsequenceCount(int[] X, int S) {
        int counter = 0;

        for (int i = 0; i < X.length; i++) {
            int[] dpSum = new int[X.length];

            dpSum[i] = X[i];

            if (X[i] == S) {
                counter++;
            }

            for (int j = i + 1; j < X.length; j++) {
                int sum = dpSum[j - 1] + X[j];

                dpSum[j] = sum;

                if ((double) sum / (j - i + 1) == S) {
                    counter++;
                }
            }
        }
        return counter;
    }

- selman

---

利用一棵树。您可以遍历数组一次，将每个元素放入一棵树中。树的每个节点应该保存值+到达该节点的值路径的平均值。最终，所有平均值为S的叶子节点的路径都是解决方案。运行时间为O(n log n)。 - Polygnome

我投票关闭此问题，因为它根据此元数据应该属于Software Recommendations.SE。 - Polygnome

值得说明的是 - 你所说的“子序列”是什么意思。通常，数组的连续片段被称为子数组，而子序列可能是不连续的。你的代码处理连续的子数组。 - MBo

这个问题的根源是什么？ - Ekesh Kumar

@Polygnome 你能详细说明一下“达到节点的值路径平均数”是什么意思吗？ 对于每个节点都应该有许多到达自身的路径，因此会有许多均值，对吧？ - selman

Answer 1

有一个技巧可以获得一个O(n)算法：

average = (A[i] + A[i+1] ... + A[j]) / (j - i + 1)

average * (j - i + 1) = A[i] + A[i+1]...+ A[j]

注意到由于average现在乘以等式右侧的元素数量，我们可以为每个元素仅减去一次平均值：

0 = (A[i]-average) + (A[i+1]-average) ... + (A[j]-average)

找出和为零的连续和可以通过检查前缀和来完成。对于每个最右边的元素（A[j]-average），我们想要添加我们在之前看到相同前缀和的次数。我们对前缀和0进行了调整，以便在适用时计算数组前缀的全长。

2 1 3, avg 2

2-2 = 0    ps = 0    count = 1 (1 for the full array prefix)
1-2 = -1   ps = -1
3-2 = 1    ps = 0    count = 2 (1 for index 0 and 1 for the full array prefix)

                     total = 3

- גלעד ברקן

---

有趣！ 你能否请澄清一下如何确定“我们之前看到相同的前缀和的次数”？ 前缀和可能达到数十亿级别。 我除了在我的答案中提到的哈希表以外想不到其他东西了。 - Cătălin Frâncu

@CătălinFrâncu 前缀和的数量不能大于O(n)。例如：[1, 2, -3, 2, -2, 4, -1]，前缀和为[1, 3, 0, 2, 0, 4, 3]，相同前缀值之间的连续和(0,0] 2 - 2 = 0，(3, 3] -3 + 2 - 2 + 4 - 1 = 0，都是零。每个最右边的前缀可以与我们在其左侧找到的许多相等前缀配对（对于完整长度前缀有一个特殊情况），以形成连续的和为零。我们只需要在向右迭代时计算已经看到的值为v的前缀数。我们可以使用哈希表来存储它们。 - גלעד ברקן

使用前缀和和字典来检查先前的任何总和是否已经出现，是一个非常聪明的想法。 - Tấn Nguyên

这个想法多么聪明。 - plhn

Answer 2

我将在此算法中使用基于1的索引。这感觉就像是那些可以使用这种索引方式的情况之一。

令P为部分和数组，即P[0] = 0，且P[i] = X[1] + ... + X[i]。此外，让Q[i] = P[i] - S * i。例如，

i     0   1   2   3   4   5   6   7
-----------------------------------
X         5   3   6   2   5   5   2
P     0   5   8  14  16  21  26  28
Q     0   1   0   2   0   1   2   0

“[i,j] 的平均数（包括 i 和 j）为 S”是什么意思？根据上述符号，可以写成：

(P[j] - P[i - 1]) / (j - i + 1) = S     ==>
P[j] - P[i - 1] = S * (j - i + 1)       ==>
P[j] - P[i - 1] = S * j - S * (i - 1)   ==>
P[j] - S * j = P[i - 1] - S * (i - 1)   ==>
Q[j] = Q[i - 1]

这意味着在 Q 中任何一对相等的值都对应于平均值为S的一段范围。例如，Q 中两个值为1的值对应于范围[3,6,2,5]。Q 中四个0的值会产生6个平均值为S=4的范围：[5,3]、[6,2]、[5,5,2]、[5,3,6,2]、[6,2,5,5,2] 和 [5,3,6,2,5,5,2]。

因此，以下算法的时间复杂度也是O(n log n)，与@Polygnome的评论相同，但实现起来应该要容易得多：

• 计算 Q；
• 对 Q 进行排序；
• 对于 Q 中每个批次的k个相等值，将 k * (k - 1) / 2 添加到答案中；
• 返回答案。

如果 Q 中的值的范围足够小，则可以通过使用哈希表或计数排序将其降低到O(n)。

- Cătălin Frâncu

---

我们不需要计数或任何类型的排序。无论元素范围如何，我们都可以在O(n)时间内完成此操作。 - גלעד ברקן

Answer 3

这里提供一种利用前缀和的Java解决方案，还加入了本主题讨论的反馈。

import java.util.*;

        public static int returnSubsequenceCount(int[] X, int S)
    {
        HashMap<Integer, Integer> prefixes = new HashMap<Integer, Integer>();
        int result = 0;
        int[] P = new int[X.length + 1];
        prefixes.put(0, 1);

        int[] Q = new int[X.length + 1];
        P[0] = 0;
        Q[0] = 0;

        for (int i = 1; i < X.length + 1; i++)
        {
            P[i] = P[i - 1] + X[i - 1];
            Q[i] = P[i] - S * i;

            if (!prefixes.containsKey(Q[i]))
            {
                prefixes.put(Q[i], 1);
            }
            else
            {
                Integer temp=prefixes.get(Q[i]);
                temp++;
                prefixes.put(Q[i],temp);

            }

        }

        for (Map.Entry<Integer, Integer> entry : prefixes.entrySet())
        {
            int value = entry.getValue();
            result += value * (value - 1) / 2;
        }

        return result;
    }

- jmvcollaborator

Answer 4

我发现Cătălin Frâncu的解释非常清晰简洁。以下是带有测试的代码（Scala）：

object T3 {
  def numOfSubArraysWithMean(a: Array[Int], s: Int): Int =
    a.lazyZip(LazyList from 1)
      .foldLeft((0, List(0))) { case ((pi, q), (ai, i)) =>
        val pi2 = pi + ai
        val qi = pi2 - s * i
        (pi2, qi :: q)
      }._2
      .groupMapReduce(identity)(_=>1)(_+_)
      .withFilter { case (_, v) => v > 1 }
      .map { case (_, v) => v * (v - 1) / 2 }
      .sum
}

class T3Spec extends AnyFunSpec with Matchers {
  import T3._
  def a[A: ClassTag](aa: A*): Array[A] = aa.toArray
  it("a") {
    val data = Seq(
      (a(5,3,6,2,5,5,2),4) -> 8,
      (a(5,3,6,2,5),4) -> 4,
      (a(5,3,6,2,5,3),4) -> 7,
      (a(5,3,6,2),4) -> 3,
      (a(5,3,6,2,4),4) -> 6,
      (a(2,1,3,7,2,2,1,3),2) -> 9,
      (a(0,8,1,7,2,6,3,5),4) -> 10,
      (a(2,2,3,4,1,1),2) -> 4,
      (a(2,2,2,2,2,2),2) -> 21,
      (a(2,2,2,2),2) -> 10,
    )
    for {
      ((a, s), r) <- data
    } numOfSubArraysWithMean(a, s) shouldEqual r
  }
}

在大多数情况下，N与NlogN并不重要。但是N^2是一个问题。

- Alexey Rykhalskiy

Answer 5

Kotlin版本

fun solution(A: IntArray, S: Int): Int {
    return A.asSequence()
        .runningFold(0) { acc, i -> acc + i }
        .map { ((it % S) + S) % S }
        .groupingBy { it }
        .eachCount()
        .values
        .sumOf { it * (it - 1) / 2 }
}

- spyr01d

---

你的回答可以通过提供更多支持信息来改进。请编辑以添加进一步的细节，例如引用或文档，以便他人可以确认你的答案是正确的。您可以在帮助中心中找到有关如何编写良好答案的更多信息。 - Community

Answer 6

这是如何在Python中实现的，关键是要生成幂集

def mean_subsequence(array: list, mean: int) -> int:
result = 0
powerset = [[]]

for num in array:
    len_ = len(powerset)
    for i in range(len_):
        subsequent = powerset[i]
        #check the mean
        tmp = [num] + subsequent
        if tmp != [] and sum(tmp) / len(tmp) == mean:
            result += 1
        #update power set
        powerset.appened([num] + subsequent)
        
return result

- Anggi Permana Harianja