fork download 
  1. # Đếm số đường đi không hướng độ dài k trên cây heap (n nodes).
  2. # n: int (<=1e18), k: nonnegative int
  3. # Trả về: số đường (int)
  4.  
  5. def count_pairs(n: int, k: int) -> int:
  6.     if k == 0:
  7.         return n  # mỗi đỉnh là một đường độ dài 0
  8.  
  9.     # helper: số nodes ở "độ sâu h" trong subtree gốc tại node c (c >= 1)
  10.     # tức là số x in [c*2^h, c*2^h + 2^h -1] ∩ [1,n]
  11.     def depth_count(c: int, h: int) -> int:
  12.         if c <= 0:
  13.             return 0
  14.         start = c << h
  15.         if start > n:
  16.             return 0
  17.         end = start + (1 << h) - 1
  18.         if end <= n:
  19.             return 1 << h
  20.         return n - start + 1
  21.  
  22.     # helper: số chẵn trong đoạn [L,R] (L,R ints, L<=R)
  23.     def count_even(L: int, R: int) -> int:
  24.         if R < L:
  25.             return 0
  26.         first = L if (L % 2 == 0) else (L + 1)
  27.         if first > R:
  28.             return 0
  29.         return ((R - first) // 2) + 1
  30.  
  31.     # Precompute for h in [0..k-1]:
  32.     # M_h = floor(n / 2^h)
  33.     # partial_size_h = n - M_h * 2^h + 1 (if M_h >= 1) else 0
  34.     M = []
  35.     partial = []
  36.     pow2 = [1 << h for h in range(k)]  # safe since k small in practice
  37.     for h in range(k):
  38.         twoh = 1 << h
  39.         Mh = n // twoh
  40.         M.append(Mh)
  41.         if Mh >= 1:
  42.             partial.append(n - Mh * twoh + 1)
  43.         else:
  44.             partial.append(0)
  45.  
  46.     # sum S1 = sum_a (depth(left, k-1) + depth(right, k-1))
  47.     # left children are even indices c=2a, right are odd c=2a+1
  48.     # we compute sum over even c of depth_count(c,k-1) and sum over odd similarly
  49.     def sum_by_parity(h: int):
  50.         Mh = M[h]
  51.         if Mh == 0:
  52.             return (0, 0)
  53.         full = 1 << h
  54.         # evens among full part (1..Mh-1)
  55.         even_full_cnt = (Mh - 1) // 2
  56.         odd_full_cnt = (Mh - 1) - even_full_cnt
  57.         even_sum = even_full_cnt * full + (partial[h] if (Mh % 2 == 0) else 0)
  58.         odd_sum  = odd_full_cnt * full  + (partial[h] if (Mh % 2 == 1) else 0)
  59.         return (even_sum, odd_sum)
  60.  
  61.     even_km1, odd_km1 = sum_by_parity(k-1)
  62.     S1 = even_km1 + odd_km1
  63.  
  64.     # S2 = sum_{l=1..k-1} sum_{a} depth(left, l-1) * depth(right, k-l-1)
  65.     # = sum_{l=1..k-1} sum_{even c} f(c) * g(c+1), where f(c)=depth_count(c, l-1), g(c+1)=depth_count(c+1, k-l-1)
  66.     S2 = 0
  67.     for l in range(1, k):
  68.         h1 = l - 1
  69.         h2 = k - l - 1
  70.         M1 = M[h1]
  71.         M2 = M[h2]
  72.         if M1 == 0 or M2 == 0:
  73.             continue
  74.         # valid c range (even c) must satisfy c <= M1 and c+1 <= M2  => c <= min(M1, M2-1)
  75.         C = min(M1, M2 - 1)
  76.         if C < 2:
  77.             continue
  78.         # range [2..C] with even c only
  79.         total_even_positions = count_even(2, C)
  80.         full_prod = (1 << h1) * (1 << h2)
  81.  
  82.         # exclude special positions where f or g is partial:
  83.         specials = {}
  84.         # position c = M1 (if in range) => f(c) is partial (otherwise full)
  85.         if 2 <= M1 <= C and (M1 % 2 == 0):
  86.             specials[M1] = ('f_partial', partial[h1])
  87.         # position c = M2 - 1 (if in range) => g(c+1) is partial
  88.         c2 = M2 - 1
  89.         if 2 <= c2 <= C and (c2 % 2 == 0):
  90.             if c2 in specials:
  91.                 # both partial
  92.                 specials[c2] = ('both_partial', (partial[h1], partial[h2]))
  93.             else:
  94.                 specials[c2] = ('g_partial', partial[h2])
  95.  
  96.         # count how many even positions are 'regular full' (neither special)
  97.         num_special_positions = len(specials)
  98.         regular_count = total_even_positions - num_special_positions
  99.         if regular_count > 0:
  100.             S2 += regular_count * full_prod
  101.  
  102.         # add contributions of special positions
  103.         for cpos, info in specials.items():
  104.             if info[0] == 'f_partial':
  105.                 fval = info[1]
  106.                 gval = (1 << h2) if (cpos + 1 != M2) else partial[h2]  # but cpos+1==M2 can't happen here because we marked only M1; safe to compute directly:
  107.                 # safer compute explicitly:
  108.                 gval = depth_count(cpos + 1, h2)
  109.                 S2 += fval * gval
  110.             elif info[0] == 'g_partial':
  111.                 gval = info[1]
  112.                 fval = depth_count(cpos, h1)
  113.                 S2 += fval * gval
  114.             else:  # both_partial
  115.                 p1, p2 = info[1]
  116.                 S2 += p1 * p2
  117.  
  118.     total = S1 + S2
  119.     return total
  120.  
Success #stdin #stdout 0.06s 14084KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/X4yMQh
language:
Python 3 (python  3.12)
created:
12 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!