Subarray Sum Equals K

+10

Se te da un array de números enteros nums y un número entero k. Cuenta los subarrays cuyos elementos suman exactamente k. Un subarray es una secuencia de uno o más elementos contiguos. Dos subarrays cuentan por separado cuando empiezan o terminan en posiciones distintas, incluso si contienen los mismos valores. Los valores pueden ser negativos o cero.

Función

subarraySum(nums: integer-array, k: integer) → integer

numsinteger-array: el arreglo de enteros, que puede contener valores negativos y ceros
kinteger: la suma que debe alcanzar un subarreglo para que se cuente
Devuelveinteger: el número de subarreglos cuyos elementos suman k

Restricciones

1 ≤ nums.length ≤ 2 × 10⁴
-1000 ≤ nums[i] ≤ 1000
-10⁷ ≤ k ≤ 10⁷
Un arreglo de esta longitud tiene como máximo 200,010,000 subarreglos, por lo que la respuesta cabe en un entero con signo de 32 bits.

Ejemplos

Entrada: nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
Salida: 4
Explicación: Cuatro secuencias suman 7: [3, 4], [1, 3, 3], [3, 3, 1] y [3, 4, -7, 1, 3, 3]. En la última, el -7 cancela el 3 y el 4, y la suma vuelve a subir a 7 más adelante, así que una secuencia puede coincidir incluso después de que su suma haya superado k.

Entrada: nums = [1, -1, 0]k = 0
Salida: 3
Explicación: Tres subarreglos suman 0: [1, -1], [0] y el arreglo completo [1, -1, 0]. El segmento [-1, 0] suma -1, así que no cuenta.

Entrada: nums = [2, 2, 2]k = 4
Salida: 2
Explicación: La secuencia [2, 2] en los índices 0 y 1 y la secuencia [2, 2] en los índices 1 y 2 contienen los mismos valores, pero están en posiciones diferentes, así que ambas cuentan. La suma de todo el arreglo es 6.

+17 pruebas ocultas al enviar

Para ir más allá

¿Cómo cambiarías la solución para que devuelva la longitud del subarreglo más largo cuya suma sea k, manteniendo un tiempo de ejecución de O(n)?

Pistas

Ábrelas de una en una. Cada una revela un poco más.

0 de 3 abiertas

Solución

Un arreglo de n números tiene n(n+1)/2 subarreglos, aproximadamente 2 × 10^8 cuando n = 2 × 10^4, así que sumarlos uno por uno es demasiado lento. Los valores negativos también descartan una ventana deslizante: la suma de una ventana puede disminuir y volver a aumentar, así que no hay ninguna regla que indique cuándo reducirla. La idea que resuelve el problema es escribir la suma de cada subarreglo como la diferencia entre dos sumas de prefijos. Contar los subarreglos que terminan en el elemento actual y suman k significa entonces contar las sumas de prefijos anteriores que son iguales a la suma de prefijos actual menos k, y un mapa hash lo resuelve en una sola pasada.

Cada inicio con un total acumulado

Correcto, pero no termina con las pruebas más grandes

Intuición

Cada subarreglo tiene un índice inicial start y un índice final end. Si recorres cada par y compruebas su suma, encuentras cada subarreglo exactamente una vez, así que el recuento es correcto.

No necesitas un tercer bucle para sumar cada subarreglo. Fija start, luego mueve end un paso a la derecha cada vez y suma nums[end] a un total acumulado. El total siempre contiene la suma de los elementos desde start hasta end, así que cada subarreglo requiere una suma y una comparación.

No te detengas cuando el total alcance o supere k. Un valor negativo posterior puede hacer que vuelva a bajar: en el primer ejemplo, el total desde el índice 0 pasa por 3, 7, 0, 1, 4, 7, así que ese inicio tiene una segunda coincidencia en el índice 5.

El costo es la cantidad de pares. Con n = 2 × 10^4 hay aproximadamente 2 × 10^8, una cantidad manejable para C, pero demasiado lenta para Python, Ruby o R.

Algoritmo

Establece count en 0.
Para cada start de 0 a n-1, establece total en 0.
Para cada end desde start hasta n-1, suma nums[end] a total.
Si total es igual a k, suma 1 a count y continúa de cualquier manera.
Devuelve count.

def subarraySum(nums, k):
    n = len(nums)
    count = 0
    for start in range(n):
        total = 0
        # Grow the subarray that begins at start, one element at a time
        for end in range(start, n):
            total += nums[end]
            if total == k:
                count += 1
    return count

Sumas de prefijos con un mapa de conteo

Intuición

Sea prefix[j] la suma de los primeros j elementos, con prefix[0] = 0 para el prefijo vacío. El subarreglo desde el índice i hasta el índice j-1 suma prefix[j] - prefix[i]. Así que un subarreglo que termina en el elemento actual suma k exactamente cuando una suma de prefijo anterior es igual a la suma de prefijo actual menos k. Cada uno de esos prefijos anteriores marca dónde empieza un subarreglo coincidente.

Recorre el arreglo una vez. Mantén la suma de prefijo acumulada y un mapa hash seen que asocie cada suma de prefijo con la cantidad de veces que ha aparecido. En cada elemento, primero suma seen[prefix - k] al contador y luego registra el prefijo actual. Buscar antes de registrar evita que un subarreglo esté vacío: con k = 0, si registraras primero, la suma de prefijo actual coincidiría consigo misma.

Considera el primer ejemplo con k = 7. Las sumas de prefijo son 0, 3, 7, 0, 1, 4, 7, 8, 4. Cuando la suma de prefijo llega a 7 después del índice 1, el mapa contiene un 0, que da [3, 4]. Cuando llega a 7 de nuevo después del índice 5, el mapa contiene dos 0, el prefijo vacío y el prefijo después del -7, que dan [3, 4, -7, 1, 3, 3] y [1, 3, 3] a la vez. En 8, después del índice 6, el mapa contiene un 1, que da [3, 3, 1]. Eso da un total de 4.

Empezar el mapa con un 0 visto una vez es lo que cuenta los subarreglos que comienzan en el índice 0. Es importante usar un mapa de conteos en lugar de un conjunto porque la misma suma de prefijo puede repetirse, y cada copia inicia un subarreglo distinto. Cada elemento requiere una búsqueda y una actualización, así que el tiempo es O(n), y el mapa contiene como máximo n+1 claves.

Algoritmo

Crea un mapa seen con seen[0] = 1 y establece prefix y count en 0.
Para cada elemento, súmalo a prefix.
Suma seen[prefix - k] a count, leyendo como 0 una clave inexistente.
Suma 1 a seen[prefix].
Devuelve count.

def subarraySum(nums, k):
    # seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
    seen = {0: 1}
    prefix = 0
    count = 0
    for num in nums:
        prefix += num
        # Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
        count += seen.get(prefix - k, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

Errores comunes y casos límite

La mayoría de las respuestas incorrectas se deben a tratar la entrada como si todos los valores fueran positivos, o al orden de las dos operaciones con el mapa.

Una ventana deslizante que se reduce cuando la suma supera k falla con valores negativos. En el primer ejemplo devuelve 2 en lugar de 4: la ventana mantiene su extremo izquierdo en el índice 0 hasta que la suma supera 7 en el índice 6, así que nunca prueba [1, 3, 3] ni [3, 3, 1].
Omitir seen[0] = 1 hace que se pierdan todos los subarreglos que empiezan en el índice 0. Para nums = [5] y k = 5, devuelve 0 en lugar de 1.
Registrar el prefijo actual antes de la búsqueda cuenta los subarreglos vacíos cuando k es 0. Para [1, -1, 0], devuelve 6 en lugar de 3.
Usar un conjunto de sumas de prefijos en lugar de un mapa de conteos da un resultado menor cuando hay repeticiones. Para [0, 0, 0] y k = 0, la respuesta es 6, porque cada copia anterior de la misma suma de prefijo inicia un subarreglo diferente.
En la fuerza bruta, salir del bucle interno cuando el total supera k es incorrecto por la misma razón que con la ventana deslizante.

Preguntas frecuentes4

¿Cuál es la complejidad temporal de Subarray Sum Equals K?

La solución con suma de prefijos y mapa hash se ejecuta en tiempo O(n) y usa O(n) espacio adicional: una pasada, con una búsqueda y una actualización por elemento. Comprobar cada subarreglo con un total acumulado requiere tiempo O(n²), y sumar cada subarreglo desde cero requiere O(n³).

¿Por qué no funciona una ventana deslizante para la suma de subarreglo igual a K?

Una ventana deslizante se basa en que la suma aumente cuando la ventana crece y disminuya cuando se encoge, lo cual solo se cumple cuando todos los valores son positivos. Con valores negativos, una ventana cuya suma ya es demasiado grande todavía puede convertirse en una coincidencia al crecer más, así que ninguna regla te indica cuándo mover el borde izquierdo. Si todos los valores fueran positivos, una ventana deslizante lo resolvería en tiempo O(n) y espacio O(1).

¿Por qué el mapa hash empieza con 0 asociado a 1?

Esa entrada representa el prefijo vacío antes del primer elemento, cuya suma es 0. La suma de un subarreglo que comienza en el índice 0 es igual a la suma del prefijo actual menos ese prefijo vacío, así que, sin esa entrada, esos subarreglos nunca se cuentan. Para nums = [5] y k = 5, la búsqueda de 5 - 5 = 0 encuentra esa entrada y devuelve 1.

¿Se puede resolver Suma de subarreglo igual a K en O(1) de espacio adicional?

No con el método de una sola pasada. Para contar las coincidencias que terminan en un elemento, necesitas saber qué sumas de prefijos aparecieron antes de él, y puede haber hasta n+1 distintas. Sin el mapa, vuelves al total acumulado O(n²). Cuando todos los valores son positivos, una ventana deslizante cuenta los subarreglos en O(n) de tiempo y O(1) de espacio.

Python

def subarraySum(nums, k):
    # Escribe el código aquí

Casos de prueba

Caso 1

Caso 2

Caso 3

Entrada

nums = [3, 4, -7, 1, 3, 3, 1, -4]
k = 7

Esperado