Subarray Sum Equals K

+10

Você recebe um array de números inteiros nums e um número inteiro k. Conte os subarrays cujos elementos somam exatamente k. Um subarray é uma sequência de um ou mais elementos adjacentes. Dois subarrays são contados separadamente quando começam ou terminam em posições diferentes, mesmo que contenham os mesmos valores. Os valores podem ser negativos ou zero.

Função

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

numsinteger-array: o array de números inteiros, que pode conter valores negativos e zeros
kinteger: a soma que um subarray deve atingir para ser contado
Retornainteger: o número de subarrays cujos elementos somam k

Restrições

1 ≤ nums.length ≤ 2 × 10⁴
-1000 ≤ nums[i] ≤ 1000
-10⁷ ≤ k ≤ 10⁷
Um array desse tamanho tem, no máximo, 200,010,000 subarrays, então a resposta cabe em um inteiro com sinal de 32 bits.

Exemplos

Entrada: nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
Saída: 4
Explicação: Quatro sequências somam 7: [3, 4], [1, 3, 3], [3, 3, 1] e [3, 4, -7, 1, 3, 3]. Na última, o -7 cancela o 3 e o 4, e a soma volta a subir até 7 mais adiante, então uma sequência pode corresponder mesmo depois que sua soma ultrapassa k.

Entrada: nums = [1, -1, 0]k = 0
Saída: 3
Explicação: Três subarrays somam 0: [1, -1], [0] e o array inteiro [1, -1, 0]. O trecho [-1, 0] soma -1, então não conta.

Entrada: nums = [2, 2, 2]k = 4
Saída: 2
Explicação: A sequência [2, 2] nos índices 0 e 1 e a sequência [2, 2] nos índices 1 e 2 contêm os mesmos valores, mas estão em posições diferentes, então ambas são contabilizadas. A soma de todo o array é 6.

+17 testes ocultos ao enviar

Para ir além

Como você alteraria a solução para retornar o comprimento do subarray mais longo cuja soma seja k, ainda em tempo O(n)?

Dicas

Abra uma de cada vez. Cada uma revela um pouco mais.

0 de 3 abertas

Solução

Um array com n números tem n(n+1)/2 subarrays, cerca de 2 × 10^8 quando n = 2 × 10^4, então somar cada um deles é lento demais. Os valores negativos também descartam uma janela deslizante: a soma de uma janela pode diminuir e voltar a aumentar, portanto nenhuma regra indica quando reduzi-la. A ideia que resolve o problema é escrever a soma de cada subarray como a diferença entre duas somas de prefixo. Contar os subarrays que terminam no elemento atual e cuja soma é k significa, então, contar as somas de prefixo anteriores iguais à soma atual menos k, e um mapa de dispersão responde a isso em uma única passagem.

Cada início com um total acumulado

Correta, mas não termina nos maiores testes

Intuição

Cada subarray tem um índice inicial start e um índice final end. Se você visitar cada par e verificar sua soma, encontrará cada subarray exatamente uma vez, então a contagem estará correta.

Você não precisa de um terceiro loop para somar cada subarray. Fixe start, depois avance end uma posição por vez e adicione nums[end] a um total acumulado. O total sempre contém a soma dos elementos de start a end, então cada subarray custa uma adição e uma comparação.

Não pare quando o total alcançar ou ultrapassar k. Um valor negativo posterior pode reduzi-lo novamente: no primeiro exemplo, o total a partir do índice 0 passa por 3, 7, 0, 1, 4, 7, então esse início tem uma segunda correspondência no índice 5.

O custo é o número de pares. Com n = 2 × 10^4, há cerca de 2 × 10^8 pares, o que é tranquilo em C, mas lento demais para Python, Ruby ou R.

Algoritmo

Defina count como 0.
Para cada start de 0 a n-1, defina total como 0.
Para cada end de start a n-1, some nums[end] a total.
Se total for igual a k, adicione 1 a count e continue de qualquer forma.
Retorne 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

Somas prefixas com um mapa de contagem

Intuição

Seja prefix[j] a soma dos primeiros j elementos, com prefix[0] = 0 para o prefixo vazio. O subarray do índice i ao índice j-1 soma prefix[j] - prefix[i]. Portanto, um subarray que termina no elemento atual soma k exatamente quando uma soma de prefixo anterior é igual à soma de prefixo atual menos k. Cada soma de prefixo anterior desse tipo marca onde um subarray correspondente começa.

Percorra o array uma vez. Mantenha a soma de prefixo acumulada e um mapa hash seen que associa cada soma de prefixo ao número de vezes que ela apareceu. Em cada elemento, primeiro adicione seen[prefix - k] à contagem e, em seguida, registre o prefixo atual. Consultar antes de registrar impede que um subarray seja vazio: com k = 0, registrar primeiro faria a soma de prefixo atual corresponder a si mesma.

Considere o primeiro exemplo com k = 7. As somas de prefixo são 0, 3, 7, 0, 1, 4, 7, 8, 4. Quando o prefixo chega a 7 após o índice 1, o mapa contém um 0, que corresponde a [3, 4]. Quando chega a 7 novamente após o índice 5, o mapa contém dois 0s, o prefixo vazio e o prefixo após o -7, que correspondem a [3, 4, -7, 1, 3, 3] e [1, 3, 3] de uma só vez. Em 8, após o índice 6, o mapa contém um 1, que corresponde a [3, 3, 1]. Isso totaliza 4.

Iniciar o mapa com 0 registrado uma vez é o que conta os subarrays que começam no índice 0. Usar um mapa de contagens, em vez de um conjunto, é importante porque a mesma soma de prefixo pode se repetir, e cada ocorrência inicia um subarray diferente. Cada elemento exige uma consulta e uma atualização, então o tempo é O(n), e o mapa contém no máximo n+1 chaves.

Algoritmo

Crie um mapa seen com seen[0] = 1 e defina prefix e count como 0.
Para cada elemento, adicione-o a prefix.
Adicione seen[prefix - k] a count, considerando como 0 uma chave ausente.
Adicione 1 a seen[prefix].
Retorne 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

Armadilhas e casos extremos

A maioria das respostas erradas vem de tratar a entrada como se todos os valores fossem positivos ou da ordem das duas operações do mapa.

Uma janela deslizante que diminui quando a soma ultrapassa k falha com valores negativos. No primeiro exemplo, ela retorna 2 em vez de 4: a janela mantém sua extremidade esquerda no índice 0 até que a soma ultrapasse 7 no índice 6, então nunca tenta [1, 3, 3] ou [3, 3, 1].
Omitir seen[0] = 1 faz com que todos os subarrays que começam no índice 0 não sejam contabilizados. Para nums = [5] e k = 5, retorna 0 em vez de 1.
Registrar o prefixo atual antes da busca contabiliza subarrays vazios quando k é 0. Para [1, -1, 0], retorna 6 em vez de 3.
Usar um conjunto de somas de prefixo no lugar de um mapa de contagens subcontabiliza as repetições. Para [0, 0, 0] e k = 0, a resposta é 6, porque cada ocorrência anterior da mesma soma de prefixo inicia um subarray diferente.
Na força bruta, interromper o loop interno quando o total ultrapassa k está errado pelo mesmo motivo que a janela deslizante.

Perguntas frequentes4

Qual é a complexidade de tempo de Subarray Sum Equals K?

A solução com soma de prefixos e mapa hash é executada em tempo O(n) e usa O(n) de espaço extra: uma passagem, com uma consulta e uma atualização por elemento. Verificar cada subarray com um total acumulado leva tempo O(n²), e somar cada subarray desde o início leva O(n³).

Por que uma janela deslizante não funciona para Subarray Sum Equals K?

Uma janela deslizante depende de a soma aumentar quando a janela cresce e diminuir quando ela encolhe, o que só acontece quando todos os valores são positivos. Com valores negativos, uma janela cuja soma já é grande demais ainda pode se tornar uma correspondência depois de crescer mais, então nenhuma regra indica quando mover a borda esquerda. Se todos os valores fossem positivos, uma janela deslizante resolveria isso em tempo O(n) e espaço O(1).

Por que o mapa de hash começa com 0 mapeado para 1?

Essa entrada representa o prefixo vazio antes do primeiro elemento, cuja soma é 0. Um subarray que começa no índice 0 tem como soma a soma prefixada atual menos esse prefixo vazio; sem essa entrada, esses subarrays nunca são contabilizados. Para nums = [5] e k = 5, a busca por 5 - 5 = 0 encontra essa entrada e retorna 1.

É possível resolver Subarray Sum Equals K com O(1) de espaço extra?

Não com o método de uma única passagem. Para contar as correspondências que terminam em um elemento, você precisa saber quais somas de prefixo vieram antes dele, e pode haver até n+1 diferentes. Sem o mapa, você volta ao total acumulado O(n²). Quando todos os valores são positivos, uma janela deslizante conta os subarrays em O(n) tempo e O(1) espaço.

Python

def subarraySum(nums, k):
    # Escreva o código aqui

Casos de teste

Caso 1

Caso 2

Caso 3

Entrada

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

Esperado