Count pairs (A, B) such that A has X and B has Y number of set bits and A+B = C

Given two numbers x, y which denotes the number of set bits. Also given is a number C. The task is to print the number of ways in which we can form two numbers A and B such that A has x number of set bits and B has y number of set bits and A+B = C.

Examples:

Input: X = 1, Y = 1, C = 3 
Output: 2 
So two possible ways are (A = 2 and B = 1) and (A = 1 and B = 2)

Input: X = 2, Y = 2, C = 20 
Output: 3 


Approach: The above problem can be solved using bitmask DP.

  1. Initialise a 4-D DP array of size 64 * 64 * 64 * 2 as 10^18 has a maximum of 64 set bits with -1
  2. The first state of the DP array stores the number of bits traversed in C from right. The second state stores the number of set-bits used out of X and third state stores the number of set bits used out of Y. The fourth state is the carry bit which refers to the carry generated when we perform an addition operation.
  3. The recurrence will have 4 possibilities. We start from the rightmost bit position.
  4. If the bit position at C is 1, then there are four possibilities to get 1 at that index.
    • If the carry is 0, then we can use 1 bit out of X and 0 bit out of Y, or the vice-versa which generates no carry for next step.
    • If the carry is 1, then we can use 1 set bit’s from each which generates a carry for the next step else we use no set bits from X and Y which generates no carry.
  5. If the bit position at C is 0, then there are four possibilities to get 0 at that bit position.
    • If the carry is 1, then we can use 1 set bit out of X and 0 bit out of Y or the vice versa which generates a carry of 1 for the next step.
    • If the carry is 0, then we can use 1 and 1 out of X and Y respectively, which generates a carry of 1 for next step. We can also use no set bits, which generates no carry for the next step.
  6. Summation of all possibilites is stored in dp[third][seta][setb][carry] to avoid re-visiting same states.

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the above approach
#include <bits/stdc++.h>
using namespace std;
  
// Initial DP array
int dp[64][64][64][2];
  
// Recursive function to generate
// all combinations of bits
int func(int third, int seta, int setb,
         int carry, int number)
{
  
    // if the state has already been visited
    if (dp[third][seta][setb][carry] != -1)
        return dp[third][seta][setb][carry];
  
    // find if C has no more set bits on left
    int shift = (number >> third);
  
    // if no set bits are left for C
    // and there are no set bits for A and B
    // and the carry is 0, then
    // this combination is possible
    if (shift == 0 and seta == 0 and setb == 0 and carry == 0)
        return 1;
  
    // if no set bits are left for C and
    // requirement of set bits for A and B have exceeded
    if (shift == 0 or seta < 0 or setb < 0)
        return 0;
  
    // Find if the bit is 1 or 0 at
    // third index to the left
    int mask = shift & 1;
  
    dp[third][seta][setb][carry] = 0;
  
    // carry = 1 and bit set = 1
    if ((mask) && carry) {
  
        // since carry is 1, and we need 1 at C's bit position
        // we can use 0 and 0
        // or 1 and 1 at A and B bit position
        dp[third][seta][setb][carry]
                += func(third + 1, seta, setb, 0, number)
                + func(third + 1, seta - 1, setb - 1, 1, number);
    }
  
    // carry = 0 and bit set = 1
    else if (mask && !carry) {
  
        // since carry is 0, and we need 1 at C's bit position
        // we can use 1 and 0
        // or 0 and 1 at A and B bit position
        dp[third][seta][setb][carry] 
                += func(third + 1, seta - 1, setb, 0, number)
                + func(third + 1, seta, setb - 1, 0, number);
    }
  
    // carry = 1 and bit set = 0
    else if (!mask && carry) {
  
        // since carry is 1, and we need 0 at C's bit position
        // we can use 1 and 0
        // or 0 and 1 at A and B bit position
        dp[third][seta][setb][carry] 
                  += func(third + 1, seta - 1, setb, 1, number)
                  + func(third + 1, seta, setb - 1, 1, number);
    }
  
    // carry = 0 and bit set = 0
    else if (!mask && !carry) {
  
        // since carry is 0, and we need 0 at C's bit position
        // we can use 0 and 0
        // or 1 and 1 at A and B bit position
        dp[third][seta][setb][carry] 
                  += func(third + 1, seta, setb, 0, number)
                  + func(third + 1, seta - 1, setb - 1, 1, number);
    }
  
    return dp[third][seta][setb][carry];
}
  
// Function to count ways
int possibleSwaps(int a, int b, int c)
{
  
    memset(dp, -1, sizeof(dp));
  
    // function call that returns the
    // answer
    int ans = func(0, a, b, 0, c);
  
    return ans;
}
  
// Driver Code
int main()
{
  
    int x = 2, y = 2, c = 20;
  
    cout << possibleSwaps(x, y, c);
  
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the above approach 
import java.util.*;
class GfG { 
  
// Initial DP array 
static int dp[][][][] = new int[64][64][64][2]; 
  
// Recursive function to generate 
// all combinations of bits 
static int func(int third, int seta, int setb, 
                int carry, int number) 
  
    // if the state has already been visited 
    if (dp[third][seta][setb][carry] != -1
        return dp[third][seta][setb][carry]; 
  
    // find if C has no more set bits on left 
    int shift = (number >> third); 
  
    // if no set bits are left for C 
    // and there are no set bits for A and B 
    // and the carry is 0, then 
    // this combination is possible 
    if (shift == 0 && seta == 0 && setb == 0 && carry == 0
        return 1
  
    // if no set bits are left for C and 
    // requirement of set bits for A and B have exceeded 
    if (shift == 0 || seta < 0 || setb < 0
        return 0
  
    // Find if the bit is 1 or 0 at 
    // third index to the left 
    int mask = shift & 1
  
    dp[third][seta][setb][carry] = 0
  
    // carry = 1 and bit set = 1 
    if ((mask == 1) && carry == 1) { 
  
        // since carry is 1, and we need 1 at C's bit position 
        // we can use 0 and 0 
        // or 1 and 1 at A and B bit position 
        dp[third][seta][setb][carry] 
                            += func(third + 1, seta, setb, 0, number) 
                            + func(third + 1, seta - 1, setb - 1, 1, number); 
    
  
    // carry = 0 and bit set = 1 
    else if (mask == 1 && carry == 0) { 
  
        // since carry is 0, and we need 1 at C's bit position 
        // we can use 1 and 0 
        // or 0 and 1 at A and B bit position 
        dp[third][seta][setb][carry] 
                                    += func(third + 1, seta - 1, setb, 0, number) 
                                    + func(third + 1, seta, setb - 1, 0, number); 
    
  
    // carry = 1 and bit set = 0 
    else if (mask == 0 && carry == 1) { 
  
        // since carry is 1, and we need 0 at C's bit position 
        // we can use 1 and 0 
        // or 0 and 1 at A and B bit position 
        dp[third][seta][setb][carry] += func(third + 1, seta - 1, setb, 1, number) 
                                    + func(third + 1, seta, setb - 1, 1, number); 
    
  
    // carry = 0 and bit set = 0 
    else if (mask == 0 && carry == 0) { 
  
        // since carry is 0, and we need 0 at C's bit position 
        // we can use 0 and 0 
        // or 1 and 1 at A and B bit position 
        dp[third][seta][setb][carry] += func(third + 1, seta, setb, 0, number) 
                                    + func(third + 1, seta - 1, setb - 1, 1, number); 
    
  
    return dp[third][seta][setb][carry]; 
  
// Function to count ways 
static int possibleSwaps(int a, int b, int c) 
    for(int q = 0; q < 64; q++)
    {
        for(int r = 0; r < 64; r++)
        {
            for(int p = 0; p < 64; p++)
            {
                for(int d = 0; d < 2; d++)
                {
                    dp[q][r][p][d] = -1;
                }
            }
        }
    }
      
  
    // function call that returns the 
    // answer 
    int ans = func(0, a, b, 0, c); 
  
    return ans; 
  
// Driver Code 
public static void main(String[] args) 
  
    int x = 2, y = 2, c = 20
  
    System.out.println(possibleSwaps(x, y, c)); 
}

chevron_right


Output:

3


My Personal Notes arrow_drop_up

Striver(underscore)79 at Codechef and codeforces D

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.



Improved By : prerna saini