LeetCode 2859.计算K置位下标对应元素的和

给你一个下标从 0 开始的整数数组 nums 和一个整数 k 。

请你用整数形式返回 nums 中的特定元素之 和 ，这些特定元素满足：其对应下标的二进制表示中恰存在 k 个置位。

整数的二进制表示中的 1 就是这个整数的 置位 。

例如，21 的二进制表示为 10101 ，其中有 3 个置位。

示例 1：

输入：nums = [5,10,1,5,2], k = 1
输出：13
解释：下标的二进制表示是：
0 = 0002
1 = 0012
2 = 0102
3 = 0112
4 = 1002
下标 1、2 和 4 在其二进制表示中都存在 k = 1 个置位。
因此，答案为 nums[1] + nums[2] + nums[4] = 13 。
示例 2：

输入：nums = [4,3,2,1], k = 2
输出：1
解释：下标的二进制表示是：
0 = 002
1 = 012
2 = 102
3 = 112
只有下标 3 的二进制表示中存在 k = 2 个置位。
因此，答案为 nums[3] = 1 。

提示：

1 <= nums.length <= 1000
1 <= nums[i] <= 105
0 <= k <= 10

法一：直接模拟：

class Solution {
public:
    int sumIndicesWithKSetBits(vector<int>& nums, int k) {
        int ans = 0;
        for (int i = 0; i < nums.size(); ++i)
        {
            if (count1bit(i) == k)
            {
                ans += nums[i];
            }
        }

        return ans;
    }

private:
    int count1bit(int n)
    {
        int ans = 0;
        while (n)
        {
            ++ans;
            n = n & (n - 1);
        }
        return ans;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

如果nums的长度为n，nums中元素为m，则此算法时间复杂度为O(nlogm)，空间复杂度为O(1)。

法二：我们可以优化求一个二进制数中1的个数的过程，题目中规定了数组的大小最大为1000，可以用10位二进制数表示，我们将其表示为abcdefghij，我们要求的实际是a+b+c+d+e+f+g+h+i+j，可以用分治法求和：
第一步：求0a0c0e0g0i+0b0d0f0h0j，由于0a+0b一定不会产生进位，最多为2位，因此我们可以将ab的和看做一个数字(ab)，即结果是(ab)(cd)(ef)(gh)(ij)，此时(ab)+(cd)+(ef)+(gh)+(ij)的结果等于a+b+c+d+e+f+g+h+i+j。

第二步：重复上一步，即(ab)+(cdef)+(ghij)的结果等于a+b+c+d+e+f+g+h+i+j。

第三步：此时只有3项了，可以直接求，即算出(ab)、(cdef)、(ghij)的值并相加即可：

class Solution {
public:
    int sumIndicesWithKSetBits(vector<int>& nums, int k) {
        int ans = 0;
        for (int i = 0; i < nums.size(); ++i)
        {
            if (count1bit(i) == k)
            {
                ans += nums[i];
            }
        }

        return ans;
    }

private:
    int count1bit(int n)
    {
        n = (0b0101010101 & n) + ((0b1010101010 & n) >> 1);
        n = ((0b0011001100 & n) >> 2) + (0b1100110011 & n);
        return (n >> 8) + ((n >> 4) & 0b1111) + n & 0b1111;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

如果nums的大小为n，nums最大为m，则此算法时间复杂度为O(nloglogm)，因为m有logm个二进制位（如上例中的10），而我们计算这logm个二进制位中的1时又用了分治算法，因此再取一层对数；空间复杂度为O(1)。

法三：使用内置函数__builtin_popcount：

class Solution {
public:
    int sumIndicesWithKSetBits(vector<int>& nums, int k) {
        int ans = 0;
        for (int i = 0; i < nums.size(); ++i)
        {
            if (__builtin_popcount(i) == k)
            {
                ans += nums[i];
            }
        }

        return ans;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

__builtin_popcount函数的时间复杂度为O(1)，如果nums的大小为n，此算法时间复杂度为O(n)，空间复杂度为O(1)。

法四：从最小的k置位数开始，计算所有的k置位数，这样只需遍历符合k置位数遍：

class Solution {
public:
    int sumIndicesWithKSetBits(vector<int>& nums, int k) {
        if (k == 0)
        {
            return nums[0];
        }

        int ans = 0;
        int index = (1 << k) - 1;
        while (index < nums.size())
        {
            ans += nums[index];

            // 每次只加lowbit，保证1的位数只会减少
            index += index & (-index);
            int bit1Num = __builtin_popcount(index);
            // 由于每次都加lowbit导致的1的位数减少，因此低位一定是0
            if (bit1Num < k)
            {
                index |= (1 << k - bit1Num) - 1;
            }
        }

        return ans;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

如果有n个K置位数，此算法时间复杂度为O(n)，空间复杂度为O(1)。

法五：gosper’s hack，类似法四，gosper’s hack的作用是，给定一个K置位数后，求出比该K置位数大的最小的K置位数，相当于将所有K置位数排序后，给定一个K置位数，获取排序后的下一个K置位数。步骤为，找到当前K置位数最右边的01，然后将其变为10，再将10右面的1放到最右面，以下是一组3置位数，可供参考：
00111
01011
01101
01110
10011
10101
10110
11001
11010
11100

对于Gosper’s Hack的实现，直接看代码：

class Solution {
public:
    int sumIndicesWithKSetBits(vector<int>& nums, int k) {
        if (k == 0)
        {
            return nums[0];
        }

        int ans = 0;
        int index = (1 << k) - 1;
        while (index < nums.size())
        {
            ans += nums[index];

            // 如果index是10110，则lowbit是10
            int lowbit = index & (-index);
            // temp是11000，temp计算出的是index的下一个K置位数的高位部分
            int temp = index + lowbit;
        
            // temp和index的异或结果为01到右面的最后一个1是1，其他是0
            // 即异或结果为01110，异或结果有3个1，是因为01...1（x个1）在10110中有3位
            // 接下来需要把index中01后面的1全部移动到最右边，即右移2（01）加上01...10...0中的最后面的0的个数位
            // 即__builtin_ctz(lowbit) + 2位，再与temp相与，就是下一个K置位数
            index = ((temp ^ index) >> __builtin_ctz(lowbit) + 2) | temp;
            // or index = (((temp ^ index) >> 2) / lowbit) | temp;
        }

        return ans;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

如果有n个K置位数，此算法时间复杂度为O(n)，空间复杂度为O(1)。