Recently, my roommate took the Huawei software written test, and after I got the questions, I also tried to solve them. I found the second question to be very clever, so I decided to write a blog post to record it.
Problem Description#
Driver Xiao Wang is at the warehouse to pick up goods. The warehouse manager has already packed the goods into boxes and arranged them in a row. The boxes are numbered starting from 1. The goods stored in each box form a set M={M1, M2, ..., Mn}. The truck can transport a maximum of K goods at a time. In order to maximize efficiency, Xiao Wang wants to select a multiple of K goods from the set and transport them in batches. The warehouse manager requires Xiao Wang to choose boxes with consecutive numbers. For example, he can choose boxes 1, 2, and 3, but he cannot choose boxes 2, 4, and 6.
Please help Xiao Wang calculate how many ways there are to select a multiple of K goods for transportation.
Input Description#
The number of boxes is N, and the maximum number of goods that can be transported by the truck at a time is K.
The number of goods stored in each box: M1, M2, M3, M4, ...
The input consists of two lines in the following format:
N K
The values of N and K are both in the range [1, 100000].
The number of goods stored in each box, Mi, is also in the range [1, 100000].
Analysis#
Since this is a consecutive sequence problem, I initially considered using a sliding window, but after analysis, I found that it is not feasible.
The general method of a sliding window is to first expand the window, and when a certain condition occurs that makes the window unable to satisfy the condition, shrink the window. However, in this problem, it is different: it is not possible to shrink the window based on a local condition. In other words, it is not possible to guarantee that the window can be expanded to satisfy a multiple of K just because the current window does not satisfy a multiple of K.
Then I thought of using dynamic programming to calculate the sum of intervals and then traverse to check if each interval meets the condition. Experience tells us that this is correct, but dynamic programming requires a two-dimensional array to mark the interval sums, with a space complexity of O(N^2), and it also needs to traverse this two-dimensional array, with a time complexity of O(N^2). In the case where N is in the range [1, 100000], such complexity is unacceptable.
Enumerating through all possibilities? Although it reduces the space complexity, the time complexity is much higher than dynamic programming, making it even more unacceptable.
When in doubt, I looked at some solutions. I randomly found this solution, which provides a solution with a time complexity of O(n+k) and gave me great inspiration.
The basic idea of this solution is based on the following principles:
- Take the difference between the sum of goods from 0 to j and the sum of goods from 0 to i, and the difference is the sum of goods from i to j. If this sum can be evenly divided by K, that is, the sum of goods from 0 to j modulo K is equal to the sum of goods from 0 to i modulo K, then the interval satisfies the condition.
The mathematical representation of this conclusion is:
(Sum[j]-Sum[i])%k=0 => Sum[j] = n*k + Sum[i] => Sum[j]%k = Sum[i]%k
- When multiple positions a, b, c... satisfy Sum[a]%k=Sum[b]%k=Sum[c]%k=..., let the number of positions be t, then any two positions among these positions can form an interval that satisfies the condition. Therefore, the number of intervals that satisfy the condition at this time is C(t,2).
C(t,2) is a mathematical combination, which expands to t*(t-1)/2.
The specific steps are as follows:
-
Use an array M of length n to store the number of goods in the warehouse. After inputting, create another array Sum of length n, where Sum[i] represents the sum of goods from M[0] to M[i]. At this time, Sum[j]-Sum[i] represents the total number of goods in the interval from goods i to goods j.
-
Create an array count of length k, and traverse the Sum array. For any index x between 0 and n, increment count[Sum[x]%k] by 1.
-
Traverse the count array. For any index y between 0 and k, if count[y]>1, then result = result + C(count[y],2). Finally, return the result.
In the actual code representation, we can omit the M and Sum arrays and perform related operations on the count array while inputting. In this case, we only need to process n inputs, traverse an array of length k, with a time complexity of O(n+k), and only use an array count of length k, with a space complexity of O(k), which is the optimal solution for this problem.
Code#
Here is my Golang code:
package main
import "fmt"
func main() {
var (
n, k int
tmp int
current int
count []int
result int
)
fmt.Scanln(&n, &k)
count = make([]int, k)
// Input the list of goods and accumulate in real-time
for i := 0; i < n; i++ {
fmt.Scan(&tmp)
current = (current + tmp) % k
count[current] += 1
}
for _, v := range count {
if v > 1 {
result += v * (v - 1) / 2
}
}
fmt.Println(result)
}