Given an array A
, you can repeatedly delete any contiguous subarrays with length k
that sum to tar
from it. Output whether A
can be made empty through this process.
A
is an integer array, k
is a positive integer, tar
is an integer.
For example, A=[1,2,3,4];k=2;tar=5
Then you can delete [2,3]
from A
so that it becomes [1,4]
. Then you delete [1,4]
from A
; it becomes empty. Therefore, the algorithm should output True
.
Currently I have found a O(n^2/k*comb(n/k,k))
algorithm. Is there a better one?
Firstly, use dp, find whether A[i:j]
can be empty, then enumerate all elements of the last deleted subarray, in time O(comb((j-i)/k, k))
,
An example code of python3, when k
is 3
:
from functools import lru_cache
from itertools import accumulate
def solve(A, k, tar):
# only works when k = 3
assert k==3
presum = list(accumulate(A, initial=0))
@lru_cache(None)
def judge(i, j):
"""
find if s[i:j] can be full erased
"""
l,remainder = divmod(j-i,k)
if remainder != 0 or presum[j]-presum[i] != tar*l:
return False
if l<=1:
return True
# enumerate the last 3 elements to be erased
# # find triplets that sum to tar
for a1 in range(i, j, 3):
for a2 in range(a1+1, j, 3):
for a3 in range(a2+1, j, 3):
if A[a1]+A[a2]+A[a3]==tar:
prv = i
flag = True
for nxt in a1,a2,a3,j:
if not judge(prv, nxt):
flag = False
break
prv = nxt+1
if flag == True:
return True
return False
return judge(0,len(A))
print(solve([1,2,3,4,8,0],3,9))
I also wonder if there is any specific algorithm when k
is 3
and 2
.
A testcase that when k=3
greedy doesn't work:
A=[5,5,5,1,9,7,3,0,10, 5,5,5, 10,0,3,7,9,1,5,5,5];tar=15;k=3
For problems where we're dealing with a sequence of operations, you can often get a dynamic programming solution by parameterizing on the first or last operation you will perform.
For our array A = [a_0, a_1, ... a_{n-1}]
, let's guess what the last move will be. It must be to remove k
elements summing to our target
, and those elements [x_0, x_1, ... x_{k-1}]
form a subsequence of A which partitions the other elements of A
into at most k+1
other subarrays. Before performing the last operation, the other subarrays have already been completely and individually deleted by some operations, and the subproblem on each of those subarrays is completely analogous to our original problem. This an example of optimal substructure.
Now that the framework is set up, let's talk about details. There's lots of easy edge cases: if k
is 1, or if the length of the array isn't divisible by k
, or if the sum of the array is incorrect, that we should check first. I'll also be assuming that target
and the entries of A
are nonnegative, but this can be modified with small additions to the code without affecting the runtime complexity.
What variables do our subproblems need to keep track of? There's the left and right bounds of our current array, as well as the number and sum of all the outer partition elements we've already taken. We don't know yet which k
partition elements will be in the last operation-- they could be spread out over the entire original array.
We'll be walking through the array from left to right, at each point testing whether we can use the current element for that outer partition, or whether to process the next k, or 2k, or 3k, ... elements together as a contiguous subarray and continue on after that. We can filter out most subarrays from consideration: a subarray S is potentially solvable if its length is some multiple of k, m*k
, and its sum is m*target
. We can use prefix sums to test this quickly.
Python
def solve(arr: List[int], k: int, target: int) -> bool:
"""Decide whether arr is completely removable
Given an array of nonnegative integers arr, a positive integer length k,
and a nonnegative target sum, decide whether we can remove all elements
of arr by repeatedly removing k consecutive elements summing to target.
"""
assert k > 0
assert target >= 0
# This section deals with any easy edge cases
if len(arr) % k != 0:
return False
if k == 1:
return all(x == target for x in arr)
prefixes = list(itertools.accumulate(arr, initial=0))
if target == 0: # Assumes target >= 0 and all values >= 0
return prefixes[-1] == 0
if prefixes[-1] != target * (len(arr) // k):
return False
if k == len(arr):
return True
reverse_index = collections.defaultdict(list)
starts_to_ends = collections.defaultdict(list) # Candidate subarrays
for right, x in enumerate(prefixes):
remainder = x % target
for left in reverse_index[remainder]:
if (right - left) % k == 0:
if (right - left) // k == (x - prefixes[left]) // target: # Valid sum
starts_to_ends[left].append(right-1)
reverse_index[remainder].append(right)
@functools.lru_cache(None)
def can_solve(left: int, right, outer_need: int, outer_sum_need: int) -> bool:
elements_remaining = right - left + 1
# We've reached the end of our array
if elements_remaining == 0:
return outer_need == outer_sum_need == 0
# All elements must go to the outer partition
if elements_remaining == outer_need:
return prefixes[right + 1] - prefixes[left] == outer_sum_need
# We've used k elements for the outer partition, but
# their sum is too small
if outer_need == 0 and outer_sum_need != 0:
return False
# Test whether we can split into two subproblems
for poss_ending in starts_to_ends[left]:
if poss_ending >= right:
break
if (can_solve(left, poss_ending, 0, 0)
and can_solve(poss_ending + 1, right, outer_need, outer_sum_need)):
return True
# We are just starting a subproblem: Initialize outer partition as empty
# and try using current element in outer partition
if outer_need == 0 and outer_sum_need == 0:
if (arr[left] <= target
and can_solve(left + 1, right, k - 1, target - arr[left])):
return True
# Try using current element in outer partition
if outer_need > 0 and arr[left] <= outer_sum_need:
if can_solve(left + 1, right, outer_need - 1, outer_sum_need - arr[left]):
return True
return False
return can_solve(left=0, right=len(arr) - 1, outer_need=0, outer_sum_need=0)
This can be modified to actually output the arrays to delete in order, with some more work. The current runtime is O(n^2 * target * n/k)
(saving a factor of k because of our pre-filtering trick when the sum or length of an array is wrong). For many solvable arrays, a 'greedy' strategy will work and be much faster, so an optimization could be to try that first.
I believe (and have some empirical evidence) that the runtime can be brought down to O(n^2 * target)
, by eliminating the loop inside of the recursive function. Instead of testing all 'chunks' of the next k
, or 2k
, or 3k
elements as a separate subproblem, it appears that you only need to test the largest such chunk whose sum is valid, although I haven't found a satisfactory proof that this always gives the right answer.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With